YAM world model: tokenizer + dynamics weights, model card, media

.gitattributes

@@ -33,3 +33,10 @@ 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/abc_domain_gap.webp filter=lfs diff=lfs merge=lfs -text
+assets/yam_action_control.webp filter=lfs diff=lfs merge=lfs -text
+assets/yam_reconstruction.webp filter=lfs diff=lfs merge=lfs -text
+yam_dynamics/model/6000/variables/ocdbt.process_0/d/53984876ba74f53c7270e8cf67f8f097 filter=lfs diff=lfs merge=lfs -text
+yam_dynamics/model/6000/variables/ocdbt.process_0/d/f846ff2a604e205aa4c0a53b67f2a9b2 filter=lfs diff=lfs merge=lfs -text
+yam_tokenizer/model/12000/variables/ocdbt.process_0/d/3803b1c360f0ac7dc881ea36a077bd79 filter=lfs diff=lfs merge=lfs -text
+yam_tokenizer/model/12000/variables/ocdbt.process_0/d/5eb28d9fe5e1e389cb00731c3f12edee filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,173 @@
+---
+license: apache-2.0
+library_name: jax
+pipeline_tag: video-to-video
+tags:
+  - robotics
+  - world-model
+  - dreamer
+  - dreamer-v4
+  - yam
+  - lerobot
+  - jax
+  - flax
+datasets:
+  - Dimios45/yam-pick-duster-200
+  - Dimios45/yam-pick-duster
+  - Dimios45/yam_towel_fold
+  - Dimios45/molmo_eval_90cmtopcam
+  - Dimios45/molmo_eval_hitl
+---
+
+# YAM World Model — a Dreamer-4 style action-conditioned video world model
+
+An action-conditioned latent video world model for the **YAM** robot arm
+(single-arm 7-DoF and bimanual 14-DoF), trained from scratch on a single
+**RTX 4090**. Given a few frames of context and a stream of joint commands, it
+imagines the future.
+
+Built on the [visionary](https://github.com/james0248/visionary) Dreamer-4
+implementation, retargeted from SO-101 to YAM.
+
+<p align="center">
+  <img src="assets/yam_action_control.webp" alt="The same clip imagined under true, shuffled, and zero actions" width="800">
+</p>
+
+## The model actually obeys its actions
+
+This is the property that matters and the one that is easy to fake. A world model
+trained on a single repetitive scene learns to predict the future from pixels
+alone and quietly ignores the action input. The control is to re-run the same
+clip, same seed, and change **only** the actions:
+
+| `--action_source` | PSNR | SSIM | rollout motion | vs ground truth |
+| --- | --- | --- | --- | --- |
+| `true` | **15.03** | **0.6675** | 3.684 | 1.6x |
+| `shuffled` (another episode's actions) | 14.55 | 0.6484 | 5.311 | **2.3x** |
+| `zero` | 17.85 | 0.7459 | 1.264 | 0.55x |
+| ground truth | — | — | 2.295 | 1.0x |
+
+* `true` beats `shuffled` on both metrics.
+* Wrong actions produce **wrong motion** — the arm thrashes at 2.3x ground truth.
+* Zero actions nearly **freeze** the arm, which is correct for "no commanded
+  movement".
+
+> ⚠️ **`zero` scores highest on PSNR, and that is a metric artifact.** On a mostly
+> static scene a near-frozen prediction beats sharp-but-slightly-misaligned motion,
+> because pixel metrics reward blur. PSNR alone would have scored this model as
+> ignoring its actions. Pair the action control with a motion statistic.
+
+## Reconstruction
+
+<p align="center">
+  <img src="assets/yam_reconstruction.webp" alt="Tokenizer reconstruction against ground truth" width="800">
+</p>
+
+| stage | metric | value |
+| --- | --- | --- |
+| tokenizer | held-out PSNR | **25.6 dB** |
+| tokenizer | mse / lpips | 0.0009 / 0.0166 |
+| tokenizer | dead latent channels | **0 / 16** |
+| dynamics | rollout PSNR / SSIM | **24.82 / 0.913** |
+| dynamics | `flow_mse` (baseline 1.0) | 0.0048 |
+
+Rollout quality saturates at the tokenizer's own reconstruction ceiling — past
+that, the dynamics model is limited by what the tokenizer can decode.
+
+## Architecture
+
+Two stages, both JAX/Flax, following Dreamer 4.
+
+**Video tokenizer** — masked-autoencoder spatiotemporal transformer. 144×192
+frames, patch 16, → 96 latent tokens × 16 channels per frame, `tanh` bottleneck.
+Video only; never sees actions. 8 encoder + 8 decoder layers, dim 512. MSE + LPIPS,
+Muon optimizer, 12,000 steps.
+
+**Dynamics** — shortcut/flow-matching transformer over the frozen latents,
+conditioned on a continuous action vector through a **per-embodiment** MLP
+projection plus an embodiment embedding. 8 layers, dim 768, 24-frame context
+(4.8 s at 5 Hz). 6,000 steps, EMA 0.999 (exports carry EMA weights).
+
+Two embodiments share one video backbone: `yam_absolute_joint_7d_v1` (7-D) and
+`bi_yam_absolute_joint_14d_v1` (14-D), `max_action_dim: 14`.
+
+## Usage
+
+```python
+from predict import WorldModel   # scripts/robot/predict.py
+
+wm = WorldModel(
+    tokenizer_dir="yam_tokenizer",
+    dynamics_dir="yam_dynamics",
+    latent_stats="latent_stats.json",
+    dynamics_step=6000,
+)
+# context_frames: (T0, H, W, 3) uint8   actions: (T0+horizon, 7) normalized to [-1,1]
+frames = wm.imagine(context_frames, actions, embodiment_id=0)
+```
+
+**5.4 fps warm** on an RTX 4090 (first call ~28 s of JIT compile; 4.4 s per
+24-frame rollout thereafter). Trained at 5 Hz, so it imagines at roughly the rate
+the data was recorded — enough for single-trajectory lookahead in a 5 Hz control
+loop, not for dense multi-rollout MPC.
+
+## Training data
+
+| Corpus | Robot | Action | Episodes | Frames |
+| --- | --- | --- | --- | --- |
+| `yam-pick-duster` + `-200` | YAM | 7-D joint, **radians** | 250 | ~104k |
+| `yam_towel_fold` | bi-YAM | 14-D joint, radians | 11 | 9,468 |
+| `molmo_eval_90cmtopcam` + `_hitl` | molmoact dual-arm | 14-D joint, radians | 5 | 13,003 |
+
+266 packed records / **118,982 frames**, 0 invalid. All are LeRobot **v3.0**. Only
+verified-fixed cameras are kept; wrist cameras are dropped. Episodes are trimmed to
+their non-idle span, actions normalized q01–q99 → [-1,1], video packed at 240×320
+and decoded to 144×192 at 5 Hz.
+
+## Limitations
+
+* **Not a policy.** This is a world model — it imagines, it does not act. No
+  planning or RL is included.
+* **Never drives a real arm.** Model output has not been executed on hardware.
+  Check your robot's calibration convention before attempting it.
+* **Single fixed viewpoint per episode**, and the scene diversity is modest
+  (duster picking, towel folding, two eval sets). Expect degradation on unseen
+  scenes, lighting, or camera placements.
+* **Bimanual is the weaker embodiment** — 16 episodes against 250 for the
+  single-arm side.
+* **Rollouts drift over long horizons**, as flow-matching video models do. Results
+  here are for a 32-frame (6.4 s) horizon from 4 context frames.
+
+### A negative result worth publishing
+
+Adding 39 episodes from [`rerun/abc-130k`](https://huggingface.co/buckets/rerun/abc-130k)
+(a Rerun conversion of the gated [`XDOF/ABC-130k`](https://huggingface.co/datasets/XDOF/ABC-130k))
+more than doubled the corpus to 260,731 frames and turned the starved bimanual
+embodiment into the larger one — and **made the model worse**. The single-arm side
+*lost* its action-conditioning (`shuffled` began outscoring `true`) and the
+bimanual side never gained any.
+
+<p align="center">
+  <img src="assets/abc_domain_gap.webp" alt="Combined tokenizer: YAM 25.6 dB versus ABC-130k 19.0 dB" width="800">
+</p>
+
+Two unseparated causes: the two bimanual robots were merged into one embodiment
+(same action *convention*, different **kinematics** — identical joint angles
+produce different pixel motion on a different robot), and the combined tokenizer
+reconstructs ABC at only 19.0 dB against YAM's 25.6, a gap that *widened* with
+training as ABC's dense two-arm scenes saturated the 96-latent bottleneck.
+
+**The weights published here are the single-corpus YAM model, which passes the
+action control.**
+
+## Reproducing
+
+Full recipe, per-stage commands, and the failure modes encountered are documented
+in [`docs/yam_world_model.md`](https://github.com/Dimios45/visionary/blob/main/docs/yam_world_model.md).
+
+## Acknowledgements
+
+* [Dreamer 4](https://danijar.com/project/dreamer4) — the architecture.
+* [visionary](https://github.com/james0248/visionary) — the implementation this builds on.
+* [Dimios45](https://huggingface.co/Dimios45) — the YAM datasets.
+* [Rerun](https://huggingface.co/rerun) and the ABC-130k authors — the bimanual corpus.

latent_stats.json

@@ -0,0 +1 @@
+{"mean": [0.1605331301689148, 0.003587166778743267, 0.3335298001766205, -0.11867159605026245, -0.23948544263839722, -0.09202785789966583, 0.02745101787149906, 0.049844179302453995, -0.08993133157491684, 0.12005544453859329, 0.040098097175359726, -0.18276961147785187, -0.13814391195774078, -0.22204774618148804, -0.31143975257873535, 0.18876683712005615], "std": [0.5045977830886841, 0.4902677834033966, 0.4043131470680237, 0.43490085005760193, 0.384480744600296, 0.39208918809890747, 0.4877498149871826, 0.4555214047431946, 0.47386112809181213, 0.4218655228614807, 0.4861060678958893, 0.5178728103637695, 0.41528770327568054, 0.44784286618232727, 0.3819977939128876, 0.44247230887413025], "count": 2076864}

norm_stats.json

@@ -0,0 +1,42 @@
+{
+ "n_episodes": 229,
+ "n_frames": 18079,
+ "action_dim": 7,
+ "normalization": "q01_q99_to_[-1,1]_clipped (gripper included)",
+ "mean": [
+  0.0683,
+  1.5591,
+  1.0881,
+  0.0349,
+  -0.0719,
+  0.0434,
+  0.2146
+ ],
+ "std": [
+  0.2481,
+  0.812,
+  0.6637,
+  0.334,
+  0.1883,
+  0.2032,
+  0.3625
+ ],
+ "q01": [
+  -0.4007,
+  0.0,
+  0.0,
+  -0.7286,
+  -0.5565,
+  -0.4525,
+  0.0
+ ],
+ "q99": [
+  0.5651,
+  2.7649,
+  2.5796,
+  0.8087,
+  0.3548,
+  0.5675,
+  1.0
+ ]
+}

yam_dynamics/model/6000/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": {"config": "orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler", "variables": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1787349275567370705, "commit_timestamp_nsecs": 1787349277105453411, "custom_metadata": {}}

yam_dynamics/model/6000/config/metadata

@@ -0,0 +1 @@
+{"_target_": "visionary.models.dreamer4.dynamics.DynamicsModel", "num_layers": 8, "num_heads": 12, "num_kv_heads": 4, "model_dim": 768, "mlp_hidden_dim": 2304, "head_dim": 64, "num_registers": 2, "num_obs_tokens": 96, "num_actions": 7, "num_embodiments": 2, "max_action_dim": 14, "action_mode": "continuous", "max_step_size": 8, "context_length": 24, "temporal_layer_period": 4, "base": 10000.0, "remat_policy": null, "latent_mean": [0.1605331301689148, 0.003587166778743267, 0.3335298001766205, -0.11867159605026245, -0.23948544263839722, -0.09202785789966583, 0.02745101787149906, 0.049844179302453995, -0.08993133157491684, 0.12005544453859329, 0.040098097175359726, -0.18276961147785187, -0.13814391195774078, -0.22204774618148804, -0.31143975257873535, 0.18876683712005615], "latent_std": [0.5045977830886841, 0.4902677834033966, 0.4043131470680237, 0.43490085005760193, 0.384480744600296, 0.39208918809890747, 0.4877498149871826, 0.4555214047431946, 0.47386112809181213, 0.4218655228614807, 0.4861060678958893, 0.5178728103637695, 0.41528770327568054, 0.44784286618232727, 0.3819977939128876, 0.44247230887413025]}

yam_dynamics/model/6000/variables/_METADATA

@@ -0,0 +1 @@
+{"tree_metadata": {"('params', 'action_embedding', 'action_projection_bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_embedding", "key_type": 2}, {"key": "action_projection_bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 768]}}, "('params', 'action_embedding', 'action_projection_kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_embedding", "key_type": 2}, {"key": "action_projection_kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 14, 768]}}, "('params', 'action_embedding', 'base_token')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_embedding", "key_type": 2}, {"key": "base_token", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'observation_projection', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "observation_projection", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'observation_projection', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "observation_projection", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 768]}}, "('params', 'prediction_projection', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "prediction_projection", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('params', 'prediction_projection', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "prediction_projection", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 16]}}, "('params', 'register_tokens')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "register_tokens", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 768]}}, "('params', 'shortcut_embedding', 'signal_embedding', 'embedding')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "shortcut_embedding", "key_type": 2}, {"key": "signal_embedding", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [129, 384]}}, "('params', 'shortcut_embedding', 'step_embedding', 'embedding')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "shortcut_embedding", "key_type": 2}, {"key": "step_embedding", "key_type": 2}, {"key": "embedding", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 384]}}, "('params', 'transformer', 'TransformerBlock_0', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_0', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_0', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_0', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_0', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_0', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_0', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_0', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_0', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2304, 768]}}, "('params', 'transformer', 'TransformerBlock_1', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_1', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_1', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_1', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_1', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_1', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_1', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_1', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_1', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_1', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_1', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2304, 768]}}, "('params', 'transformer', 'TransformerBlock_2', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_2', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_2', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_2', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_2', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_2', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_2', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_2', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_2', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_2', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_2', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2304, 768]}}, "('params', 'transformer', 'TransformerBlock_3', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_3', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_3', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_3', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_3', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_3', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_3', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_3', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_3', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_3', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_3', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2304, 768]}}, "('params', 'transformer', 'TransformerBlock_4', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_4', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_4', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_4', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_4', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_4', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_4', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_4', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_4', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_4', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_4', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2304, 768]}}, "('params', 'transformer', 'TransformerBlock_5', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_5', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_5', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_5', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_5', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_5', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_5', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_5', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_5', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_5', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_5', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2304, 768]}}, "('params', 'transformer', 'TransformerBlock_6', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_6', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_6', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_6', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_6', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_6', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_6', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_6', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_6', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_6', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_6', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2304, 768]}}, "('params', 'transformer', 'TransformerBlock_7', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_7', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_7', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 256]}}, "('params', 'transformer', 'TransformerBlock_7', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 768]}}, "('params', 'transformer', 'TransformerBlock_7', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_7', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'transformer', 'TransformerBlock_7', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_7', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'transformer', 'TransformerBlock_7', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_7', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 2304]}}, "('params', 'transformer', 'TransformerBlock_7', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2304, 768]}}, "('params', 'transformer', 'final_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "transformer", "key_type": 2}, {"key": "final_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

yam_dynamics/model/6000/variables/_sharding

@@ -0,0 +1 @@
+{"cGFyYW1zLm9ic2VydmF0aW9uX3Byb2plY3Rpb24uYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLm9ic2VydmF0aW9uX3Byb2plY3Rpb24ua2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmFjdGlvbl9lbWJlZGRpbmcuYWN0aW9uX3Byb2plY3Rpb25fYmlhcw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmFjdGlvbl9lbWJlZGRpbmcuYWN0aW9uX3Byb2plY3Rpb25fa2VybmVs":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmFjdGlvbl9lbWJlZGRpbmcuYmFzZV90b2tlbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnByZWRpY3Rpb25fcHJvamVjdGlvbi5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnByZWRpY3Rpb25fcHJvamVjdGlvbi5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnJlZ2lzdGVyX3Rva2Vucw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnNob3J0Y3V0X2VtYmVkZGluZy5zaWduYWxfZW1iZWRkaW5nLmVtYmVkZGluZw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnNob3J0Y3V0X2VtYmVkZGluZy5zdGVwX2VtYmVkZGluZy5lbWJlZGRpbmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMC5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMC5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMC5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMS5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMS5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMS5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMS5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMS5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMi5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMi5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMi5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMi5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMi5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMy5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMy5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMy5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMy5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfMy5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNC5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNC5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNC5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNS5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNS5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNS5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNS5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNS5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNi5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNi5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNi5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNi5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNi5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNy5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNy5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNy5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNy5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLlRyYW5zZm9ybWVyQmxvY2tfNy5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLnRyYW5zZm9ybWVyLmZpbmFsX25vcm0uc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}"}

yam_dynamics/model/6000/variables/array_metadatas/process_0

@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "params.action_embedding.action_projection_bias", "write_shape": [2, 768], "chunk_shape": [2, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_embedding.action_projection_kernel", "write_shape": [2, 14, 768], "chunk_shape": [2, 14, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_embedding.base_token", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.observation_projection.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.observation_projection.kernel", "write_shape": [16, 768], "chunk_shape": [16, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.prediction_projection.bias", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.prediction_projection.kernel", "write_shape": [768, 16], "chunk_shape": [768, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.register_tokens", "write_shape": [2, 768], "chunk_shape": [2, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.shortcut_embedding.signal_embedding.embedding", "write_shape": [129, 384], "chunk_shape": [129, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.shortcut_embedding.step_embedding.embedding", "write_shape": [8, 384], "chunk_shape": [8, 384], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_0.Attention_0.Dense_0.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_0.Attention_0.Dense_1.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_0.Attention_0.Dense_2.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_0.Attention_0.Dense_3.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_0.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_0.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_0.RMSNorm_0.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_0.RMSNorm_1.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_0.SwiGLU_0.Dense_0.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_0.SwiGLU_0.Dense_1.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_0.SwiGLU_0.Dense_2.kernel", "write_shape": [2304, 768], "chunk_shape": [2304, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_1.Attention_0.Dense_0.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_1.Attention_0.Dense_1.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_1.Attention_0.Dense_2.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_1.Attention_0.Dense_3.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_1.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_1.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_1.RMSNorm_0.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_1.RMSNorm_1.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_1.SwiGLU_0.Dense_0.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_1.SwiGLU_0.Dense_1.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_1.SwiGLU_0.Dense_2.kernel", "write_shape": [2304, 768], "chunk_shape": [2304, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_2.Attention_0.Dense_0.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_2.Attention_0.Dense_1.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_2.Attention_0.Dense_2.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_2.Attention_0.Dense_3.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_2.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_2.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_2.RMSNorm_0.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_2.RMSNorm_1.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_2.SwiGLU_0.Dense_0.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_2.SwiGLU_0.Dense_1.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_2.SwiGLU_0.Dense_2.kernel", "write_shape": [2304, 768], "chunk_shape": [2304, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_3.Attention_0.Dense_0.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_3.Attention_0.Dense_1.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_3.Attention_0.Dense_2.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_3.Attention_0.Dense_3.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_3.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_3.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_3.RMSNorm_0.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_3.RMSNorm_1.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_3.SwiGLU_0.Dense_0.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_3.SwiGLU_0.Dense_1.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_3.SwiGLU_0.Dense_2.kernel", "write_shape": [2304, 768], "chunk_shape": [2304, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_4.Attention_0.Dense_0.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_4.Attention_0.Dense_1.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_4.Attention_0.Dense_2.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_4.Attention_0.Dense_3.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_4.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_4.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_4.RMSNorm_0.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_4.RMSNorm_1.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_4.SwiGLU_0.Dense_0.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_4.SwiGLU_0.Dense_1.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_4.SwiGLU_0.Dense_2.kernel", "write_shape": [2304, 768], "chunk_shape": [2304, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_5.Attention_0.Dense_0.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_5.Attention_0.Dense_1.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_5.Attention_0.Dense_2.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_5.Attention_0.Dense_3.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_5.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_5.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_5.RMSNorm_0.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_5.RMSNorm_1.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_5.SwiGLU_0.Dense_0.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_5.SwiGLU_0.Dense_1.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_5.SwiGLU_0.Dense_2.kernel", "write_shape": [2304, 768], "chunk_shape": [2304, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_6.Attention_0.Dense_0.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_6.Attention_0.Dense_1.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_6.Attention_0.Dense_2.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_6.Attention_0.Dense_3.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_6.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_6.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_6.RMSNorm_0.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_6.RMSNorm_1.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_6.SwiGLU_0.Dense_0.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_6.SwiGLU_0.Dense_1.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_6.SwiGLU_0.Dense_2.kernel", "write_shape": [2304, 768], "chunk_shape": [2304, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_7.Attention_0.Dense_0.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_7.Attention_0.Dense_1.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_7.Attention_0.Dense_2.kernel", "write_shape": [768, 256], "chunk_shape": [768, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_7.Attention_0.Dense_3.kernel", "write_shape": [768, 768], "chunk_shape": [768, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_7.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_7.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_7.RMSNorm_0.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_7.RMSNorm_1.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_7.SwiGLU_0.Dense_0.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_7.SwiGLU_0.Dense_1.kernel", "write_shape": [768, 2304], "chunk_shape": [768, 2304], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.TransformerBlock_7.SwiGLU_0.Dense_2.kernel", "write_shape": [2304, 768], "chunk_shape": [2304, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.transformer.final_norm.scale", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}]}

yam_dynamics/model/6000/variables/ocdbt.process_0/d/53984876ba74f53c7270e8cf67f8f097

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

yam_dynamics/model/6000/variables/ocdbt.process_0/d/f846ff2a604e205aa4c0a53b67f2a9b2

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

yam_tokenizer/model/12000/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": {"config": "orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler", "variables": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1787347513656814520, "commit_timestamp_nsecs": 1787347515096183955, "custom_metadata": {}}

yam_tokenizer/model/12000/config/metadata

@@ -0,0 +1 @@
+{"_target_": "visionary.models.dreamer4.tokenizer.Tokenizer", "encoder_num_layers": 8, "encoder_num_heads": 8, "encoder_num_kv_heads": 2, "encoder_model_dim": 512, "encoder_mlp_hidden_dim": 1536, "decoder_num_layers": 8, "decoder_num_heads": 8, "decoder_num_kv_heads": 2, "decoder_model_dim": 512, "decoder_mlp_hidden_dim": 1536, "num_latents": 96, "head_dim": 64, "channel_dim": 16, "temporal_layer_period": 4, "temporal_layer_offset": 1, "resize_shape": [144, 192], "pad_width": [0, 0], "patch_size": 16, "base": 10000.0, "independent_prob": 0.2}

yam_tokenizer/model/12000/variables/_METADATA

@@ -0,0 +1 @@
+{"tree_metadata": {"('params', 'decoder', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16, 512]}}, "('params', 'decoder', 'Dense_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'decoder', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 768]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'decoder', 'SpatioTemporalTransformer_0', 'final_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "final_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'decoder', 'image_tokens')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "decoder", "key_type": 2}, {"key": "image_tokens", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [108, 512]}}, "('params', 'encoder', 'Dense_0', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768, 512]}}, "('params', 'encoder', 'Dense_1', 'bias')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [16]}}, "('params', 'encoder', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 16]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_0', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_0", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_1', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_1", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_2', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_2", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_3', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_3", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_4', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_4", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_5', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_5", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_6', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_6", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'Attention_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'Attention_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'Attention_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 128]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'Attention_0', 'Dense_3', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "Dense_3", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'Attention_0', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'Attention_0', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "Attention_0", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'RMSNorm_0', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "RMSNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'RMSNorm_1', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "RMSNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'SwiGLU_0', 'Dense_0', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'SwiGLU_0', 'Dense_1', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1536]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'TransformerBlock_7', 'SwiGLU_0', 'Dense_2', 'kernel')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "TransformerBlock_7", "key_type": 2}, {"key": "SwiGLU_0", "key_type": 2}, {"key": "Dense_2", "key_type": 2}, {"key": "kernel", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1536, 512]}}, "('params', 'encoder', 'SpatioTemporalTransformer_0', 'final_norm', 'scale')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "SpatioTemporalTransformer_0", "key_type": 2}, {"key": "final_norm", "key_type": 2}, {"key": "scale", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'encoder', 'latent_tokens')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "latent_tokens", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [96, 512]}}, "('params', 'encoder', 'mask_token')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "encoder", "key_type": 2}, {"key": "mask_token", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}}, "use_ocdbt": true, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

yam_tokenizer/model/12000/variables/_sharding

@@ -0,0 +1 @@
+{"cGFyYW1zLmRlY29kZXIuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLmZpbmFsX25vcm0uc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmRlY29kZXIuaW1hZ2VfdG9rZW5z":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuRGVuc2VfMC5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuRGVuc2VfMC5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuRGVuc2VfMS5iaWFz":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuRGVuc2VfMS5rZXJuZWw=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMC5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMS5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMi5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfMy5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNC5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNS5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNi5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5EZW5zZV8zLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5BdHRlbnRpb25fMC5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5STVNOb3JtXzAuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5STVNOb3JtXzEuc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5Td2lHTFVfMC5EZW5zZV8wLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5Td2lHTFVfMC5EZW5zZV8xLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLlRyYW5zZm9ybWVyQmxvY2tfNy5Td2lHTFVfMC5EZW5zZV8yLmtlcm5lbA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIuU3BhdGlvVGVtcG9yYWxUcmFuc2Zvcm1lcl8wLmZpbmFsX25vcm0uc2NhbGU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIubGF0ZW50X3Rva2Vucw==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}","cGFyYW1zLmVuY29kZXIubWFza190b2tlbg==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1], \"axis_names\": [\"data\"], \"axis_types\": [\"AxisType.Auto\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [{\"id\": 0}]}}"}

yam_tokenizer/model/12000/variables/array_metadatas/process_0

@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "params.decoder.Dense_0.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.Dense_0.kernel", "write_shape": [16, 512], "chunk_shape": [16, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.Dense_1.bias", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.Dense_1.kernel", "write_shape": [512, 768], "chunk_shape": [512, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_0.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_0.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_0.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_0.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_0.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_0.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_0.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_0.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_0.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_0.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_0.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_1.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_1.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_1.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_1.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_1.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_1.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_1.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_1.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_1.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_1.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_1.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_2.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_2.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_2.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_2.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_2.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_2.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_2.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_2.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_2.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_2.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_2.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_3.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_3.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_3.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_3.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_3.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_3.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_3.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_3.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_3.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_3.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_3.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_4.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_4.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_4.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_4.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_4.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_4.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_4.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_4.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_4.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_4.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_4.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_5.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_5.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_5.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_5.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_5.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_5.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_5.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_5.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_5.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_5.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_5.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_6.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_6.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_6.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_6.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_6.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_6.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_6.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_6.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_6.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_6.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_6.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_7.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_7.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_7.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_7.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_7.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_7.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_7.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_7.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_7.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_7.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.TransformerBlock_7.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.SpatioTemporalTransformer_0.final_norm.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.decoder.image_tokens", "write_shape": [108, 512], "chunk_shape": [108, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.Dense_0.bias", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.Dense_0.kernel", "write_shape": [768, 512], "chunk_shape": [768, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.Dense_1.bias", "write_shape": [16], "chunk_shape": [16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.Dense_1.kernel", "write_shape": [512, 16], "chunk_shape": [512, 16], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_0.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_0.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_0.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_0.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_0.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_0.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_0.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_0.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_0.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_0.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_0.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_1.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_1.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_1.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_1.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_1.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_1.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_1.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_1.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_1.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_1.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_1.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_2.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_2.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_2.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_2.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_2.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_2.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_2.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_2.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_2.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_2.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_2.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_3.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_3.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_3.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_3.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_3.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_3.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_3.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_3.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_3.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_3.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_3.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_4.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_4.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_4.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_4.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_4.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_4.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_4.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_4.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_4.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_4.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_4.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_5.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_5.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_5.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_5.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_5.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_5.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_5.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_5.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_5.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_5.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_5.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_6.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_6.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_6.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_6.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_6.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_6.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_6.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_6.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_6.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_6.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_6.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_7.Attention_0.Dense_0.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_7.Attention_0.Dense_1.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_7.Attention_0.Dense_2.kernel", "write_shape": [512, 128], "chunk_shape": [512, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_7.Attention_0.Dense_3.kernel", "write_shape": [512, 512], "chunk_shape": [512, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_7.Attention_0.RMSNorm_0.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_7.Attention_0.RMSNorm_1.scale", "write_shape": [64], "chunk_shape": [64], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_7.RMSNorm_0.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_7.RMSNorm_1.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_7.SwiGLU_0.Dense_0.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_7.SwiGLU_0.Dense_1.kernel", "write_shape": [512, 1536], "chunk_shape": [512, 1536], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.TransformerBlock_7.SwiGLU_0.Dense_2.kernel", "write_shape": [1536, 512], "chunk_shape": [1536, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.SpatioTemporalTransformer_0.final_norm.scale", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.latent_tokens", "write_shape": [96, 512], "chunk_shape": [96, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.encoder.mask_token", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}]}

yam_tokenizer/model/12000/variables/ocdbt.process_0/d/3803b1c360f0ac7dc881ea36a077bd79

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

yam_tokenizer/model/12000/variables/ocdbt.process_0/d/5eb28d9fe5e1e389cb00731c3f12edee

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

yam_tokenizer/preprocessor/12000/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": "orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1787344362030082243, "commit_timestamp_nsecs": 1787344362038858944, "custom_metadata": {}}

yam_tokenizer/preprocessor/12000/metadata

@@ -0,0 +1 @@
+{"resize_shape": [144, 192], "pad_width": [0, 0], "patch_size": 16}