Publish simulation-only IMPRINT models

.gitattributes

@@ -33,3 +33,57 @@ 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
+simulation/rmbench/full_attention/stage2/step-100000/params/d/4e59861641970db34a401178d94abccb filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/8db5b2133a434ad1141ea8812c66d055 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/2d36b6774904b98fd2859c445da7ccf9 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/59fc128f2c4b5113c11703c155e01798 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/1a7cb12cfafe1b40b4192cc5c88b6261 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/3a9a6d211c80055a23cce98ac2846aaa filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/3079b18dc746bda4c96ea07b0cc49164 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/078a5e4801dabdb2efd36a5b376784f1 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/16efd2bcc545a7d01205a2fc1a735fae filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/7dc0a88b6d199ccbb016750c98012cf1 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/recurrent/stage2/step-100000/params/d/cbf2fef1100ba6c460caea5d0023e183 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fa7cea5bceef1307d13d9a61ea91a35e filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/3a109d42459c2150f8b988ce2442e2b5 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fb91f39ce5a047cdcaebc0bad36caad6 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/b7327531c46c5cc244cc7981e9f4da90 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/054944d2cae72ec6dcde8844c294284c filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/88cd20aa3cd4faa14d8123fbe93a44b9 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/recurrent/stage2/step-100000/params/ocdbt.process_0/d/fa71f45e6411d633e6827c04a989a92f filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/d/c4e769d6e9f190c84d9d4da645a382e9 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/bfffdc95d40b690cb5c008b7d725f065 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/48fe4c11a9c01f4660d334ae365a7b60 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/2983ddb0f602eacdef62886c8414a0b9 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6fb80018af6be2a5b502f7c4b9c7cde1 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/54b5a8bb7fc15a85fec99bcf452d39d8 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/83f0bd6ad3426d6da01c9b19642be891 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/6f7030b0a150604b069a4d80c82fbca5 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/d9edf3ee1160af4f686433b3913e0761 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/full_attention/stage2/step-150000/params/ocdbt.process_0/d/dc752638f1824862095b922c6eeece25 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/58653ce97c5c65682ffda96d132e00a9 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/3ba25031cc9c39f9d751672811335cfb filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/d/bdb8580854eaecb69c9d7240af3fe6e1 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/773bae5278dbb4b09f534065a38e6191 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/2ad4b381ab460f243b187ea893aa0075 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/e1ece44be3ce97bde62b1308aefdab19 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/af24ffd2755407dd4a54220bdf599a01 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/f3b0b08919720e731107ef7b1a373a20 filter=lfs diff=lfs merge=lfs -text
+simulation/robocasa/recurrent/stage2/step-150000/params/ocdbt.process_0/d/52e1af74bd01a88827c2cdec55bb33b2 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/pi0/stage2/step-100000/params/d/3bd5e7a25312cf2bc1713db9a1d29fa5 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/3c822efa587f8375db55af5ad6ff4923 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/5c003ca5eacb776438fd966f8fe51bdf filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/24d3ba3b910a02ad89c9e9c5ee773b7d filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/bc0fe1585dbbc9796ccd89ef1cb53d7e filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/e364328538b1fe570bf0733fb7086ce5 filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/181a1cb7822e6e4ae74e5d52fd4c3e2c filter=lfs diff=lfs merge=lfs -text
+simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/6f0b57f972221d7c08e1c6e0542b0f59 filter=lfs diff=lfs merge=lfs -text
+real_world/dust_table_100/full_attention/stage1/step-50000/checkpoint/__0_0.distcp filter=lfs diff=lfs merge=lfs -text
+real_world/dust_table_100/full_attention/stage2/step-50000/checkpoint/.metadata filter=lfs diff=lfs merge=lfs -text
+real_world/dust_table_100/full_attention/stage2/step-50000/checkpoint/__0_0.distcp filter=lfs diff=lfs merge=lfs -text
+real_world/multitask_400/full_attention/balanced/stage1/step-200000/checkpoint/__0_0.distcp filter=lfs diff=lfs merge=lfs -text
+real_world/multitask_400/full_attention/linear_reward/stage1/step-200000/checkpoint/__0_0.distcp filter=lfs diff=lfs merge=lfs -text
+real_world/multitask_400/full_attention/balanced/stage2/step-200000/checkpoint/.metadata filter=lfs diff=lfs merge=lfs -text
+real_world/multitask_400/full_attention/balanced/stage2/step-200000/checkpoint/__0_0.distcp filter=lfs diff=lfs merge=lfs -text
+real_world/multitask_400/full_attention/linear_reward/stage2/step-200000/checkpoint/.metadata filter=lfs diff=lfs merge=lfs -text
+real_world/multitask_400/full_attention/linear_reward/stage2/step-200000/checkpoint/__0_0.distcp filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,37 @@
+---
+license: other
+tags:
+  - robotics
+  - vision-language-action
+  - memory
+  - imprint
+---
+
+# IMPRINT Models
+
+This repository contains the public model checkpoints for IMPRINT.
+
+## Layout
+
+| Path | Benchmark | Memory encoder |
+| --- | --- | --- |
+| `simulation/rmbench/full_attention/` | RMBench | Full attention |
+| `simulation/rmbench/recurrent/` | RMBench | Recurrent gated delta |
+| `simulation/rmbench/pi0/` | RMBench | Full attention + native Pi0 |
+| `simulation/rmbench/molmoact2/` | RMBench | Full attention + MolmoAct2 |
+| `simulation/robocasa/full_attention/` | Customized RoboCasa | Full attention |
+| `simulation/robocasa/recurrent/` | Customized RoboCasa | Recurrent gated delta |
+
+The Pi0.5 directories preserve the files from their original release
+repositories. The MolmoAct2 directory contains a model-only conversion of the
+paper checkpoint. Optimizer, dataloader, scheduler, and trainer state are
+omitted. See
+`release_manifest.json` for immutable provenance, file counts, byte counts,
+and checksums.
+
+Code, configuration, and evaluation instructions are published in the
+`Amazon-FAR/imprint` GitHub repository.
+
+The checkpoints remain subject to their applicable upstream model licenses
+and terms, including Gemma/PaliGemma where used. The repository metadata does
+not relicense upstream weights.

release_manifest.json

@@ -0,0 +1,94 @@
+{
+  "format_version": 1,
+  "release": "imprint-v1",
+  "migration_revisions": {
+    "datasets": "refs/pr/1",
+    "models": "refs/pr/1"
+  },
+  "destinations": {
+    "datasets": {
+      "repo_id": "Amazon-FAR/imprint_datasets",
+      "repo_type": "dataset"
+    },
+    "models": {
+      "repo_id": "Amazon-FAR/imprint_models",
+      "repo_type": "model"
+    }
+  },
+  "artifacts": {
+    "simulation/rmbench/450": {
+      "source": "yilin-wu/imprint-rmbench-450",
+      "source_revision": "0875f60a1a0fb79751cde5c502b5e1923a157512",
+      "files": 350,
+      "bytes": 411250782303,
+      "destination": "Amazon-FAR/imprint_datasets"
+    },
+    "simulation/robocasa/600": {
+      "source": "yilin-wu/imprint-robocasa-600",
+      "source_revision": "979c8e896855f393d560ef43aabc4e68bc4e6248",
+      "files": 465,
+      "bytes": 590046837296,
+      "destination": "Amazon-FAR/imprint_datasets"
+    },
+    "simulation/robocasa/environment": {
+      "source": "yilin-wu/imprint-robocasa-environment",
+      "source_revision": "86b8c2f042ad03776b0d23d3cd582e62b7fedacf",
+      "files": 6,
+      "bytes": 41746512,
+      "destination": "Amazon-FAR/imprint_datasets"
+    },
+    "simulation/rmbench/full_attention": {
+      "source": "yilin-wu/imprint-rmbench-full-attention",
+      "source_revision": "7be585a21fa03b9d114d771c0bc67f5e244e7b06",
+      "files": 29,
+      "bytes": 11163847579,
+      "destination": "Amazon-FAR/imprint_models"
+    },
+    "simulation/rmbench/recurrent": {
+      "source": "yilin-wu/imprint-rmbench-recurrent",
+      "source_revision": "f4e20a5d54bdbb562769175625c0926369ab1d46",
+      "files": 31,
+      "bytes": 12406277703,
+      "destination": "Amazon-FAR/imprint_models"
+    },
+    "simulation/robocasa/full_attention": {
+      "source": "yilin-wu/imprint-robocasa-full-attention",
+      "source_revision": "ce51086f892028fb5bf669aa3e7debe1b75998cc",
+      "files": 32,
+      "bytes": 12291932686,
+      "destination": "Amazon-FAR/imprint_models"
+    },
+    "simulation/robocasa/recurrent": {
+      "source": "yilin-wu/imprint-robocasa-recurrent",
+      "source_revision": "54b23b08dc766db75827defaafe718bb18b40b73",
+      "files": 29,
+      "bytes": 12406163068,
+      "destination": "Amazon-FAR/imprint_models"
+    },
+    "simulation/rmbench/molmoact2": {
+      "source": "W&B run mit0jkx7, model-only conversion at step 100000",
+      "source_revision": "mit0jkx7:100000",
+      "files": 4,
+      "bytes": 22747370042,
+      "destination": "Amazon-FAR/imprint_models",
+      "sha256": {
+        "checkpoint/model.pt": "ae0dacfcb2a45e82d18266ff5ed49edc26c35c41e038a43065c9a5671773360c",
+        "checkpoint/config.yaml": "ed4573be786a951ccb3319a090b8443db03acb1a204c827327eb6164ecfb905d"
+      }
+    },
+    "simulation/rmbench/molmoact2_eval_assets": {
+      "source": "W&B run mit0jkx7 evaluation metadata and stored language embeddings",
+      "source_revision": "mit0jkx7:100000",
+      "files": 454,
+      "bytes": 3899744,
+      "destination": "Amazon-FAR/imprint_datasets"
+    },
+    "simulation/rmbench/pi0": {
+      "source": "W&B run s746sl7k, native Pi0 model-only checkpoint",
+      "source_revision": "s746sl7k:100000",
+      "files": 24,
+      "bytes": 11534442762,
+      "destination": "Amazon-FAR/imprint_models"
+    }
+  }
+}

simulation/rmbench/full_attention/.gitattributes

@@ -0,0 +1,44 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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
+stage2/step-100000/params/d/4e59861641970db34a401178d94abccb filter=lfs diff=lfs merge=lfs -text
+stage2/step-100000/params/ocdbt.process_0/d/078a5e4801dabdb2efd36a5b376784f1 filter=lfs diff=lfs merge=lfs -text
+stage2/step-100000/params/ocdbt.process_0/d/8db5b2133a434ad1141ea8812c66d055 filter=lfs diff=lfs merge=lfs -text
+stage2/step-100000/params/ocdbt.process_0/d/16efd2bcc545a7d01205a2fc1a735fae filter=lfs diff=lfs merge=lfs -text
+stage2/step-100000/params/ocdbt.process_0/d/59fc128f2c4b5113c11703c155e01798 filter=lfs diff=lfs merge=lfs -text
+stage2/step-100000/params/ocdbt.process_0/d/3079b18dc746bda4c96ea07b0cc49164 filter=lfs diff=lfs merge=lfs -text
+stage2/step-100000/params/ocdbt.process_0/d/3a9a6d211c80055a23cce98ac2846aaa filter=lfs diff=lfs merge=lfs -text
+stage2/step-100000/params/ocdbt.process_0/d/2d36b6774904b98fd2859c445da7ccf9 filter=lfs diff=lfs merge=lfs -text
+stage2/step-100000/params/ocdbt.process_0/d/1a7cb12cfafe1b40b4192cc5c88b6261 filter=lfs diff=lfs merge=lfs -text

simulation/rmbench/full_attention/LICENSE_GEMMA.txt

@@ -0,0 +1,77 @@
+Gemma Terms of Use 
+
+Last modified: February 21, 2024
+
+By using, reproducing, modifying, distributing, performing or displaying any portion or element of Gemma, Model Derivatives including via any Hosted Service, (each as defined below) (collectively, the "Gemma Services") or otherwise accepting the terms of this Agreement, you agree to be bound by this Agreement.
+
+Section 1: DEFINITIONS
+1.1 Definitions
+(a) "Agreement" or "Gemma Terms of Use" means these terms and conditions that govern the use, reproduction, Distribution or modification of the Gemma Services and any terms and conditions incorporated by reference.
+
+(b) "Distribution" or "Distribute" means any transmission, publication, or other sharing of Gemma or Model Derivatives to a third party, including by providing or making Gemma or its functionality available as a hosted service via API, web access, or any other electronic or remote means ("Hosted Service").
+
+(c) "Gemma" means the set of machine learning language models, trained model weights and parameters identified at ai.google.dev/gemma, regardless of the source that you obtained it from.
+
+(d) "Google" means Google LLC.
+
+(e) "Model Derivatives" means all (i) modifications to Gemma, (ii) works based on Gemma, or (iii) any other machine learning model which is created by transfer of patterns of the weights, parameters, operations, or Output of Gemma, to that model in order to cause that model to perform similarly to Gemma, including distillation methods that use intermediate data representations or methods based on the generation of synthetic data Outputs by Gemma for training that model. For clarity, Outputs are not deemed Model Derivatives.
+
+(f) "Output" means the information content output of Gemma or a Model Derivative that results from operating or otherwise using Gemma or the Model Derivative, including via a Hosted Service.
+
+1.2
+As used in this Agreement, "including" means "including without limitation".
+
+Section 2: ELIGIBILITY AND USAGE
+2.1 Eligibility
+You represent and warrant that you have the legal capacity to enter into this Agreement (including being of sufficient age of consent). If you are accessing or using any of the Gemma Services for or on behalf of a legal entity, (a) you are entering into this Agreement on behalf of yourself and that legal entity, (b) you represent and warrant that you have the authority to act on behalf of and bind that entity to this Agreement and (c) references to "you" or "your" in the remainder of this Agreement refers to both you (as an individual) and that entity.
+
+2.2 Use
+You may use, reproduce, modify, Distribute, perform or display any of the Gemma Services only in accordance with the terms of this Agreement, and must not violate (or encourage or permit anyone else to violate) any term of this Agreement.
+
+Section 3: DISTRIBUTION AND RESTRICTIONS
+3.1 Distribution and Redistribution
+You may reproduce or Distribute copies of Gemma or Model Derivatives if you meet all of the following conditions:
+
+You must include the use restrictions referenced in Section 3.2 as an enforceable provision in any agreement (e.g., license agreement, terms of use, etc.) governing the use and/or distribution of Gemma or Model Derivatives and you must provide notice to subsequent users you Distribute to that Gemma or Model Derivatives are subject to the use restrictions in Section 3.2.
+You must provide all third party recipients of Gemma or Model Derivatives a copy of this Agreement.
+You must cause any modified files to carry prominent notices stating that you modified the files.
+All Distributions (other than through a Hosted Service) must be accompanied by a "Notice" text file that contains the following notice: "Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms".
+You may add your own intellectual property statement to your modifications and, except as set forth in this Section, may provide additional or different terms and conditions for use, reproduction, or Distribution of your modifications, or for any such Model Derivatives as a whole, provided your use, reproduction, modification, Distribution, performance, and display of Gemma otherwise complies with the terms and conditions of this Agreement. Any additional or different terms and conditions you impose must not conflict with the terms of this Agreement.
+
+3.2 Use Restrictions
+You must not use any of the Gemma Services:
+
+for the restricted uses set forth in the Gemma Prohibited Use Policy at ai.google.dev/gemma/prohibited_use_policy ("Prohibited Use Policy"), which is hereby incorporated by reference into this Agreement; or
+in violation of applicable laws and regulations.
+To the maximum extent permitted by law, Google reserves the right to restrict (remotely or otherwise) usage of any of the Gemma Services that Google reasonably believes are in violation of this Agreement.
+
+3.3 Generated Output
+Google claims no rights in Outputs you generate using Gemma. You and your users are solely responsible for Outputs and their subsequent uses.
+
+Section 4: ADDITIONAL PROVISIONS
+4.1 Updates
+Google may update Gemma from time to time, and you must make reasonable efforts to use the latest version of Gemma.
+
+4.2 Trademarks
+Nothing in this Agreement grants you any rights to use Google's trademarks, trade names, logos or to otherwise suggest endorsement or misrepresent the relationship between you and Google. Google reserves any rights not expressly granted herein.
+
+4.3 DISCLAIMER OF WARRANTY
+UNLESS REQUIRED BY APPLICABLE LAW, THE GEMMA SERVICES, AND OUTPUTS, ARE PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING ANY WARRANTIES OR CONDITIONS OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING, REPRODUCING, MODIFYING, PERFORMING, DISPLAYING OR OR DISTRIBUTING ANY OF THE GEMMA SERVICES OR OUTPUTS AND ASSUME ANY AND ALL RISKS ASSOCIATED WITH YOUR USE OR DISTRIBUTION OF ANY OF THE GEMMA SERVICES OR OUTPUTS AND YOUR EXERCISE OF RIGHTS AND PERMISSIONS UNDER THIS AGREEMENT.
+
+4.4 LIMITATION OF LIABILITY
+TO THE FULLEST EXTENT PERMITTED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), PRODUCT LIABILITY, CONTRACT, OR OTHERWISE, UNLESS REQUIRED BY APPLICABLE LAW, SHALL GOOGLE OR ITS AFFILIATES BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY, CONSEQUENTIAL, OR PUNITIVE DAMAGES, OR LOST PROFITS OF ANY KIND ARISING FROM THIS AGREEMENT OR RELATED TO, ANY OF THE GEMMA SERVICES OR OUTPUTS EVEN IF GOOGLE OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
+
+4.5 Term, Termination, and Survival
+The term of this Agreement will commence upon your acceptance of this Agreement (including acceptance by your use, modification, or Distribution, reproduction, performance or display of any portion or element of the Gemma Services) and will continue in full force and effect until terminated in accordance with the terms of this Agreement. Google may terminate this Agreement if you are in breach of any term of this Agreement. Upon termination of this Agreement, you must delete and cease use and Distribution of all copies of Gemma and Model Derivatives in your possession or control. Sections 1, 2.1, 3.3, 4.2 to 4.9 shall survive the termination of this Agreement.
+
+4.6 Governing Law and Jurisdiction
+This Agreement will be governed by the laws of the State of California without regard to choice of law principles. The UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. The state and federal courts of Santa Clara County, California shall have exclusive jurisdiction of any dispute arising out of this Agreement.
+
+4.7 Severability
+If any provision of this Agreement is held to be invalid, illegal or unenforceable, the remaining provisions shall be unaffected thereby and remain valid as if such provision had not been set forth herein.
+
+4.8 Entire Agreement
+This Agreement states all the terms agreed between the parties and supersedes all other agreements between the parties as of the date of acceptance relating to its subject matter.
+
+4.9 No Waiver
+Google will not be treated as having waived any rights by not exercising (or delaying the exercise of) any rights under this Agreement.

simulation/rmbench/full_attention/NOTICE

@@ -0,0 +1,11 @@
+IMPRINT
+Copyright 2026
+
+This distribution contains modifications to the OpenPI project:
+https://github.com/Physical-Intelligence/openpi
+
+OpenPI is distributed under the Apache License 2.0. Gemma components and
+derivatives are subject to the Gemma Terms of Use in LICENSE_GEMMA.txt.
+
+Gemma is provided under and subject to the Gemma Terms of Use found at
+ai.google.dev/gemma/terms.

simulation/rmbench/full_attention/README.md

@@ -0,0 +1,36 @@
+---
+license: gemma
+library_name: jax
+tags:
+  - robotics
+  - vision-language-action
+  - memory
+  - rmbench
+  - dinov3
+---
+
+# IMPRINT RMBench Full-Attention Checkpoint
+
+Canonical Pi0.5 IMPRINT checkpoint for RMBench using the full-attention
+memory encoder.
+
+- Stage 1: 160,000 updates, W&B `ylinwu/Robocasa/j8t4ebpo`
+- Stage 2: 100,000 updates, W&B `liqhuiha/Robocasa/grruxd5p`
+- memory tokens: 256 x 1024
+- policy horizon: 50
+- visual memory: three-view pooled DINOv3 features
+- Stage 2: memory encoder finetuned; image tower and language embedding table frozen
+
+The release contains the Stage-1 encoder, the complete Stage-2 Orbax parameter
+checkpoint (including the PaliGemma/Gemma backbone), normalization statistics,
+portable history configuration, and provenance. The checkpoint is a Gemma
+model derivative and is distributed subject to the included Gemma Terms of
+Use.
+
+Use the matching `yilin-wu/imprint-rmbench-450` dataset and the evaluation
+contract in the IMPRINT GitHub repository. Evaluation requires strict causal
+actions, online recent DINOv3 features, `EVAL_CONSOLIDATE_ALL=1`, BGR policy
+inputs, and 50-step prediction/execution chunks.
+
+The matching training store and reproducibility code are publicly available
+under the IMPRINT release.

simulation/rmbench/full_attention/SHA256SUMS

@@ -0,0 +1,27 @@
+3e2c24001f9ef57bf7ec959a3658fbb49cdad113cdf394c264da9d16f9bdd132  ./LICENSE_GEMMA.txt
+5b21d7116ac42a0ab6141488cb654c2411557ac977e2276accb88ef774f28c19  ./NOTICE
+32a0a8a3422765c68c10bbd944cbb0d96d4781fab507b061f02f2a4099085a05  ./README.md
+a4c8e95585b546442350a3e3e7b88939569a69a092409fc32ef1203e9589d73b  ./config/history_config.yaml
+6569bbc55da67ba7f89a5250ae7261f192754bb27a7fc5db8b579b013452d61a  ./config/input_provenance.json
+372a92664fadef29d988c33d083c0c2594227ae0078272d80d1d31560086f994  ./config/norm_stats.json
+5ea9c49f1b975f0aa1c4d7275381be406954d7d5e9bf6c359c9ee6dbe32359a9  ./stage1/stage1_step160000.npz
+3280d8311fd54fd78143091908db4051671ef7dd6de4665e0298438e428b5cac  ./stage2/step-100000/_CHECKPOINT_METADATA
+372a92664fadef29d988c33d083c0c2594227ae0078272d80d1d31560086f994  ./stage2/step-100000/assets/robomme/norm_stats.json
+cc86ddddfcb438e841a5533b5f22f33c86d4e29c6268084313bcee24f6fb918c  ./stage2/step-100000/params/_METADATA
+f126ea3bbd65b911ba0a0e3d27dc3c92c6c4b072fe97807cefe8b947a0e8e66e  ./stage2/step-100000/params/_sharding
+f666d8e56c6199af6d45f6cc1c46765cb71cf3583647a13d67c8d56ee4b22356  ./stage2/step-100000/params/array_metadatas/process_0
+21bd96f6985b94e4a1521163ca6bf5e2008f322a58c9fd0e949ba6008b08d8a2  ./stage2/step-100000/params/d/4e59861641970db34a401178d94abccb
+69bddfd565fec7dcc3f500bca4fbe862a0615908f25363b7053e02b74319b962  ./stage2/step-100000/params/manifest.ocdbt
+91b19940b50c87336da9008ed29e69dfaed8cf1a83d3caff93d3664e5ac9c5be  ./stage2/step-100000/params/ocdbt.process_0/d/078a5e4801dabdb2efd36a5b376784f1
+3b5861efcd7b43d842b14554dbafb062a087edbe0f7f8a61b275bfa6d23a0e9a  ./stage2/step-100000/params/ocdbt.process_0/d/16efd2bcc545a7d01205a2fc1a735fae
+5f18bdcbb4fb31730dbf44b59cec875e876788e1fe3e814e9adfb69663df3da0  ./stage2/step-100000/params/ocdbt.process_0/d/1a7cb12cfafe1b40b4192cc5c88b6261
+43046316feb9752714eb4ea27d8cb758279120aeb5ed646374206af3983f4975  ./stage2/step-100000/params/ocdbt.process_0/d/2d36b6774904b98fd2859c445da7ccf9
+60681c50b5dcc37dfa61b64e6a0eb5de17b290463e3bebf82ddb39cde4323a7f  ./stage2/step-100000/params/ocdbt.process_0/d/3079b18dc746bda4c96ea07b0cc49164
+7b03262932b220351d793247960e57eadd53cd4b1e1be9cf7f54f5de35600ea6  ./stage2/step-100000/params/ocdbt.process_0/d/3a9a6d211c80055a23cce98ac2846aaa
+034c510b6b4a7b7a01bff53cd076593b5f8e993f6497015f311517b1feced6d7  ./stage2/step-100000/params/ocdbt.process_0/d/48e92482d01b438d4442cd8ccbcb5df5
+72b78b1724e3ec3b4eb9f23e9ce9986349691e72e3ad863b839533f3a96f922e  ./stage2/step-100000/params/ocdbt.process_0/d/59fc128f2c4b5113c11703c155e01798
+29c1fdaf3193fee770906cd10bd8b155c1a124a6b3008f364e4e56d41367c77d  ./stage2/step-100000/params/ocdbt.process_0/d/70e9c830e0a291cdb934ad90c5ca3dc1
+455068761755068aae533bc9f2f7c8f556736ad27128f90a4bfea3f63c9f003c  ./stage2/step-100000/params/ocdbt.process_0/d/8db5b2133a434ad1141ea8812c66d055
+7b85974983b9719ecf40991f70066431672edc8a4b870e766ea41b67ac4c8993  ./stage2/step-100000/params/ocdbt.process_0/d/9940789f52fa3ebe1b791f1d439deb33
+bcf80d7c5753d9a4e7c040f4695512be62299c2653395f56ced3797e1c1dc4f1  ./stage2/step-100000/params/ocdbt.process_0/d/a9d841920262430040bbe87d60eeeb29
+2b2be2c66199fcd1776de3aea5b3e2c601388e74766bbaf5b767b6ed02384c72  ./stage2/step-100000/params/ocdbt.process_0/manifest.ocdbt

simulation/rmbench/full_attention/config/history_config.yaml

@@ -0,0 +1,68 @@
+budget: 256
+num_views: 3
+streaming_obs_horizon: 16
+integration_type: modulation
+memory_token_dim: 1024
+representation_type: cosmo_bottleneck
+memory_feature:
+  img:
+    net: identity
+    input_dim: 768
+  pos:
+    input_dim: 768
+    hidden_dim: 768
+  state:
+    input_dim: 16
+    hidden_dim: 512
+cosmo_bottleneck_memory:
+  episode_cache_size: 450
+  visual_feature_stream: dinov3
+  feature_store_path: data/rmbench-450/stage1
+  latent_source: offline
+  n_variants: 8
+  num_query_samples: 16
+  variant: Wan2.2-TI2V-5B
+  temporal_factor: 4
+  downsample: 2
+  num_latents: 256
+  num_cross_blocks: 2
+  num_heads: 8
+  cosmos_latent_dim: 48
+  siglip_dim: 768
+  proprio_in_dim: 16
+  proprio_encoding: raw
+  action_in_dim: 16
+  action_encoding: raw
+  lang_emb_dim: 2048
+  action_tokenizer_hidden: 512
+  proprio_tokenizer_hidden: 512
+  pool_h: 2
+  pool_w: 2
+  wan_h: 256
+  wan_w: 256
+  latent_cap: 4096
+  max_recent_tokens: 768
+  max_step_tokens: 1244
+cosmo_decoder:
+  num_heads: 8
+  num_blocks: 2
+  align_dim: 768
+  align_num_patches: 192
+  ffn_inner: 8192
+  action_dim: 16
+  proprio_dim: 16
+  max_query_steps: 2048
+frame_jitter:
+  enabled: false
+  start_jitter: 3
+  stride_choices:
+    - 1
+    - 2
+    - 3
+  interval_jitter: 1
+  canonical_downsample: 2
+stream_aug:
+  enabled: false
+  gauss_sigma: 0.05
+  stream_drop_p: 0.1
+  feat_drop_p: 0.0

simulation/rmbench/full_attention/config/input_provenance.json

@@ -0,0 +1,25 @@
+{
+  "action_causality": "strictly_causal",
+  "architecture": "perceiver_v1",
+  "base_policy": "pi0.5",
+  "checkpoint_metadata_sha256": "3280d8311fd54fd78143091908db4051671ef7dd6de4665e0298438e428b5cac",
+  "checkpoint_path": "stage2/step-100000",
+  "checkpoint_step": 100000,
+  "dataset": "yilin-wu/imprint-rmbench-450",
+  "eval_consolidate_all": 1,
+  "execution_horizon": 50,
+  "freeze_variant": "image_and_language_embedding",
+  "latent_cap": 4096,
+  "max_recent_tokens": 768,
+  "max_step_tokens": 1244,
+  "memory_language_source": "live",
+  "num_latents": 256,
+  "prediction_horizon": 50,
+  "stage1_checkpoint": "stage1/stage1_step160000.npz",
+  "stage1_sha256": "5ea9c49f1b975f0aa1c4d7275381be406954d7d5e9bf6c359c9ee6dbe32359a9",
+  "stage1_wandb_run": "ylinwu/Robocasa/j8t4ebpo",
+  "stage2_auxiliary_loss_weight": 0.0,
+  "stage2_memory_encoder": "finetuned",
+  "stage2_wandb_run": "liqhuiha/Robocasa/grruxd5p",
+  "visual_feature_stream": "dinov3"
+}

simulation/rmbench/full_attention/config/norm_stats.json

@@ -0,0 +1,152 @@
+{
+  "norm_stats": {
+    "actions": {
+      "mean": [
+        -0.2884466350078583,
+        1.116854190826416,
+        0.8176121115684509,
+        -0.6286301016807556,
+        0.0027539930306375027,
+        0.3858006000518799,
+        0.0,
+        0.29012763500213623,
+        1.1434277296066284,
+        0.834911048412323,
+        -0.6164064407348633,
+        0.0033871293999254704,
+        0.1835649162530899,
+        0.0,
+        0.44485998153686523,
+        0.693847119808197
+      ],
+      "std": [
+        0.2924744486808777,
+        0.9123055338859558,
+        0.6921584606170654,
+        0.5167360901832581,
+        0.024822130799293518,
+        0.7730752229690552,
+        0.0,
+        0.3022353947162628,
+        0.9363216161727905,
+        0.7197534441947937,
+        0.5159713625907898,
+        0.042254358530044556,
+        0.4932859241962433,
+        0.0,
+        0.4852803945541382,
+        0.43509986996650696
+      ],
+      "q01": [
+        -1.1566282510757446,
+        -0.00001813583730836399,
+        -0.000035173223295714706,
+        -1.493809461593628,
+        -0.044773634523153305,
+        -2.2744300365448,
+        0.0,
+        -0.000023844564566388726,
+        -0.000019354150936123915,
+        -0.0000323197637044359,
+        -1.4539425373077393,
+        -0.11438962072134018,
+        -1.3077747821807861,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.000020988931282772683,
+        2.410673141479492,
+        2.2643818855285645,
+        0.000017526803276268765,
+        0.09094329178333282,
+        1.4207451343536377,
+        0.0,
+        0.9992868304252625,
+        2.5619561672210693,
+        2.3963286876678467,
+        0.000027463351216283627,
+        0.12538659572601318,
+        1.50703763961792,
+        0.0,
+        1.0,
+        1.0
+      ]
+    },
+    "state": {
+      "mean": [
+        -0.28619107604026794,
+        1.1170974969863892,
+        0.8187518119812012,
+        -0.6292704939842224,
+        0.002742639509961009,
+        0.38321176171302795,
+        0.0,
+        0.287190318107605,
+        1.1249585151672363,
+        0.8228866457939148,
+        -0.6072733998298645,
+        0.003468270180746913,
+        0.18374495208263397,
+        0.0,
+        0.47009503841400146,
+        0.700372040271759
+      ],
+      "std": [
+        0.28622451424598694,
+        0.9085981845855713,
+        0.6900886297225952,
+        0.5145737528800964,
+        0.024824900552630424,
+        0.7737277746200562,
+        0.0,
+        0.3019956946372986,
+        0.9365604519844055,
+        0.720547616481781,
+        0.5160077214241028,
+        0.04248249530792236,
+        0.4887017011642456,
+        0.0,
+        0.4873206615447998,
+        0.4321308135986328
+      ],
+      "q01": [
+        -0.991405725479126,
+        -0.00001813583730836399,
+        -0.00003476843994576484,
+        -1.4938021898269653,
+        -0.04476987197995186,
+        -2.274427652359009,
+        0.0,
+        -0.000023844564566388726,
+        -0.000019354150936123915,
+        -0.0000323197637044359,
+        -1.455243706703186,
+        -0.11395290493965149,
+        -1.3077747821807861,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.000020988931282772683,
+        2.410663366317749,
+        2.2642898559570312,
+        0.000016971260265563615,
+        0.09090737998485565,
+        1.4203816652297974,
+        0.0,
+        1.0001307725906372,
+        2.56237530708313,
+        2.397764205932617,
+        0.000027463351216283627,
+        0.12876388430595398,
+        1.4705802202224731,
+        0.0,
+        1.0,
+        1.0
+      ]
+    }
+  }
+}

simulation/rmbench/full_attention/stage1/stage1_step160000.npz

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

simulation/rmbench/full_attention/stage2/step-100000/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": {"assets": "openpi.training.checkpoints.CallbackHandler", "params": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "train_state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1788097695402693623, "commit_timestamp_nsecs": 1788097725293618197, "custom_metadata": {}}

simulation/rmbench/full_attention/stage2/step-100000/assets/robomme/norm_stats.json

@@ -0,0 +1,152 @@
+{
+  "norm_stats": {
+    "actions": {
+      "mean": [
+        -0.2884466350078583,
+        1.116854190826416,
+        0.8176121115684509,
+        -0.6286301016807556,
+        0.0027539930306375027,
+        0.3858006000518799,
+        0.0,
+        0.29012763500213623,
+        1.1434277296066284,
+        0.834911048412323,
+        -0.6164064407348633,
+        0.0033871293999254704,
+        0.1835649162530899,
+        0.0,
+        0.44485998153686523,
+        0.693847119808197
+      ],
+      "std": [
+        0.2924744486808777,
+        0.9123055338859558,
+        0.6921584606170654,
+        0.5167360901832581,
+        0.024822130799293518,
+        0.7730752229690552,
+        0.0,
+        0.3022353947162628,
+        0.9363216161727905,
+        0.7197534441947937,
+        0.5159713625907898,
+        0.042254358530044556,
+        0.4932859241962433,
+        0.0,
+        0.4852803945541382,
+        0.43509986996650696
+      ],
+      "q01": [
+        -1.1566282510757446,
+        -0.00001813583730836399,
+        -0.000035173223295714706,
+        -1.493809461593628,
+        -0.044773634523153305,
+        -2.2744300365448,
+        0.0,
+        -0.000023844564566388726,
+        -0.000019354150936123915,
+        -0.0000323197637044359,
+        -1.4539425373077393,
+        -0.11438962072134018,
+        -1.3077747821807861,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.000020988931282772683,
+        2.410673141479492,
+        2.2643818855285645,
+        0.000017526803276268765,
+        0.09094329178333282,
+        1.4207451343536377,
+        0.0,
+        0.9992868304252625,
+        2.5619561672210693,
+        2.3963286876678467,
+        0.000027463351216283627,
+        0.12538659572601318,
+        1.50703763961792,
+        0.0,
+        1.0,
+        1.0
+      ]
+    },
+    "state": {
+      "mean": [
+        -0.28619107604026794,
+        1.1170974969863892,
+        0.8187518119812012,
+        -0.6292704939842224,
+        0.002742639509961009,
+        0.38321176171302795,
+        0.0,
+        0.287190318107605,
+        1.1249585151672363,
+        0.8228866457939148,
+        -0.6072733998298645,
+        0.003468270180746913,
+        0.18374495208263397,
+        0.0,
+        0.47009503841400146,
+        0.700372040271759
+      ],
+      "std": [
+        0.28622451424598694,
+        0.9085981845855713,
+        0.6900886297225952,
+        0.5145737528800964,
+        0.024824900552630424,
+        0.7737277746200562,
+        0.0,
+        0.3019956946372986,
+        0.9365604519844055,
+        0.720547616481781,
+        0.5160077214241028,
+        0.04248249530792236,
+        0.4887017011642456,
+        0.0,
+        0.4873206615447998,
+        0.4321308135986328
+      ],
+      "q01": [
+        -0.991405725479126,
+        -0.00001813583730836399,
+        -0.00003476843994576484,
+        -1.4938021898269653,
+        -0.04476987197995186,
+        -2.274427652359009,
+        0.0,
+        -0.000023844564566388726,
+        -0.000019354150936123915,
+        -0.0000323197637044359,
+        -1.455243706703186,
+        -0.11395290493965149,
+        -1.3077747821807861,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.000020988931282772683,
+        2.410663366317749,
+        2.2642898559570312,
+        0.000016971260265563615,
+        0.09090737998485565,
+        1.4203816652297974,
+        0.0,
+        1.0001307725906372,
+        2.56237530708313,
+        2.397764205932617,
+        0.000027463351216283627,
+        0.12876388430595398,
+        1.4705802202224731,
+        0.0,
+        1.0,
+        1.0
+      ]
+    }
+  }
+}

simulation/rmbench/full_attention/stage2/step-100000/params/_METADATA

@@ -0,0 +1 @@
+{"tree_metadata": {"('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1152, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076, 1152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 16, 72, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'embedding', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'embedding', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [14, 14, 3, 288]}}, "('params', 'PaliGemma', 'img', 'head', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'PaliGemma', 'img', 'head', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 512]}}, "('params', 'PaliGemma', 'img', 'pos_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "pos_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 1152]}}, "('params', 'PaliGemma', 'llm', 'embedder', 'input_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "embedder", "key_type": 2}, {"key": "input_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64288, 2048]}}, "('params', 'PaliGemma', 'llm', 'final_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [768]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'kv_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "kv_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'mem_rms_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "mem_rms_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'out_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "out_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'q_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "q_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 2048, 4096]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4096, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 768]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 768]}}, "('params', 'action_in_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'action_in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 1024]}}, "('params', 'action_out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('params', 'action_out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 32]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'ffn', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'ffn', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_kv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_kv', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_q', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_q', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_s', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_s', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'ffn', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'ffn', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_kv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_kv', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_q', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_q', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_s', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_s', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'lang_film', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'lang_film', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('params', 'mem_encoder', 'latents', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "latents", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1024]}}, "('params', 'mem_encoder', 'proj_cosmo', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'proj_cosmo', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [12, 1024]}}, "('params', 'mem_encoder', 'proj_recent', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'proj_recent', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 1024]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 512]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 512]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'time_mlp_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'time_mlp_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'time_mlp_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'time_mlp_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "time_mlp_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

simulation/rmbench/full_attention/stage2/step-100000/params/_sharding

@@ -0,0 +1 @@
+{"cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcucG9zX2VtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZW1iZWRkZXIuaW5wdXRfZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm1lbV9ybXNfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm91dF9laW5zdW1fbWVtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLmt2X2VpbnN1bV9tZW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLnFfZWluc3VtX21lbS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXIudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmdhdGluZ19laW5zdW0udmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmxpbmVhci52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW1fMS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLkRlbnNlXzAua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fa3YuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fa3Yuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcy5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmZmbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmZmbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fa3YuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fa3Yuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcy5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmZmbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmZmbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhdGVudHMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX291dC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnRpbWVfbWxwX2luLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}"}

simulation/rmbench/full_attention/stage2/step-100000/params/array_metadatas/process_0

@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.scale.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.bias.value", "write_shape": [27, 1076], "chunk_shape": [27, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.kernel.value", "write_shape": [27, 1152, 1076], "chunk_shape": [27, 1152, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.kernel.value", "write_shape": [27, 1076, 1152], "chunk_shape": [27, 1076, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.kernel.value", "write_shape": [27, 16, 72, 288], "chunk_shape": [27, 16, 72, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.kernel.value", "write_shape": [14, 14, 3, 288], "chunk_shape": [14, 14, 3, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.kernel.value", "write_shape": [1152, 512], "chunk_shape": [1152, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.pos_embedding.value", "write_shape": [1, 64, 1152], "chunk_shape": [1, 64, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.embedder.input_embedding.value", "write_shape": [64288, 2048], "chunk_shape": [64288, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.bias.value", "write_shape": [768], "chunk_shape": [768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.Dense_0.kernel.value", "write_shape": [1024, 768], "chunk_shape": [1024, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.w.value", "write_shape": [18, 8, 256, 512], "chunk_shape": [18, 8, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum_1.w.value", "write_shape": [18, 8, 256, 256], "chunk_shape": [18, 8, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.w.value", "write_shape": [18, 2, 1, 512, 256], "chunk_shape": [18, 2, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum_1.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.w.value", "write_shape": [18, 8, 512, 256], "chunk_shape": [18, 8, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum_1.w.value", "write_shape": [18, 8, 256, 256], "chunk_shape": [18, 8, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.kv_einsum_mem.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.mem_rms_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.out_einsum_mem.w.value", "write_shape": [18, 4, 256, 256], "chunk_shape": [18, 4, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.q_einsum_mem.w.value", "write_shape": [18, 4, 256, 256], "chunk_shape": [18, 4, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.bias.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.kernel.value", "write_shape": [18, 1024, 512], "chunk_shape": [18, 1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum.value", "write_shape": [18, 2, 2048, 4096], "chunk_shape": [18, 2, 2048, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear.value", "write_shape": [18, 4096, 2048], "chunk_shape": [18, 4096, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.gating_einsum.value", "write_shape": [18, 2, 1024, 1024], "chunk_shape": [18, 2, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.linear.value", "write_shape": [18, 1024, 1024], "chunk_shape": [18, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.Dense_0.kernel.value", "write_shape": [18, 1024, 768], "chunk_shape": [18, 1024, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.bias.value", "write_shape": [18, 768], "chunk_shape": [18, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.Dense_0.kernel.value", "write_shape": [18, 1024, 768], "chunk_shape": [18, 1024, 768], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.kernel.value", "write_shape": [8, 1024], "chunk_shape": [8, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.kernel.value", "write_shape": [256, 32], "chunk_shape": [256, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.ffn.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.ffn.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_kv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_kv.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_q.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_q.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_s.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_s.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.ffn.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.ffn.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_kv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_kv.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_q.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_q.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_s.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_s.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.latents.value", "write_shape": [64, 1024], "chunk_shape": [64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.kernel.value", "write_shape": [12, 1024], "chunk_shape": [12, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.kernel.value", "write_shape": [192, 1024], "chunk_shape": [192, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.kernel.value", "write_shape": [4, 512], "chunk_shape": [4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.kernel.value", "write_shape": [4, 512], "chunk_shape": [4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_in.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.time_mlp_out.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}]}

simulation/rmbench/full_attention/stage2/step-100000/params/d/4e59861641970db34a401178d94abccb

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

simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/078a5e4801dabdb2efd36a5b376784f1

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

simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/16efd2bcc545a7d01205a2fc1a735fae

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

simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/1a7cb12cfafe1b40b4192cc5c88b6261

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

simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/2d36b6774904b98fd2859c445da7ccf9

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

simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/3079b18dc746bda4c96ea07b0cc49164

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

simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/3a9a6d211c80055a23cce98ac2846aaa

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

simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/59fc128f2c4b5113c11703c155e01798

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

simulation/rmbench/full_attention/stage2/step-100000/params/ocdbt.process_0/d/8db5b2133a434ad1141ea8812c66d055

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

simulation/rmbench/molmoact2/README.md

@@ -0,0 +1,22 @@
+---
+library_name: pytorch
+tags:
+  - robotics
+  - rmbench
+  - molmoact2
+  - vision-language-action
+---
+
+# IMPRINT RMBench MolmoAct2
+
+This directory contains the model-only 100k Stage-2 checkpoint used for the
+IMPRINT MolmoAct2 RMBench experiment (`mit0jkx7`). Optimizer and trainer state
+are excluded.
+
+The checkpoint uses the full-attention IMPRINT Stage-1 memory encoder with 256
+memory tokens. Evaluation uses online WAN and DINOv3 history features and the
+stored Pi0.5 language embeddings under
+`Amazon-FAR/imprint_datasets/simulation/rmbench/molmoact2_eval_assets/`.
+
+Use the `molmoact2` branch of the IMPRINT repository for training and
+closed-loop evaluation.

simulation/rmbench/molmoact2/checkpoint/config.yaml

@@ -0,0 +1,3172 @@
+run_name: molmoact2-stage2-dinov3-memory-reference-fixed-film-100k-v3
+model:
+  model_name: molmoact2
+  data_formatter:
+    prompt_templates: uber_model_v2
+    message_format: qwen3
+    system_prompt: demo_or_style_v2
+    always_start_with_space: false
+    default_inference_len: 65
+    select_answer: best
+    debug: false
+    image_last: false
+    format_message_list: null
+    p_one_message: 0.0
+    eval_system_prompt_mapping: null
+    p_choice_content_in_mc: 1.0
+    template_video_mc_questions: true
+    pointing_format: html-v2
+    points_decimal_places: 1
+    use_seperate_non_pointing_qa_style: false
+    timestamp_mode: 50-percent-seconds
+    output_timestamp_mode: seconds
+    seconds_decimal_places: 1
+    p_multi_point_all_image: 0.5
+    add_setup_tokens: true
+    add_control_tokens: true
+    use_seperate_count_without_pointing_style: false
+    sample_random_initial_point: true
+  llm:
+    d_model: 2560
+    n_heads: 32
+    n_kv_heads: 8
+    head_dim: 128
+    qkv_bias: false
+    clip_qkv: null
+    n_layers: 36
+    mlp_ratio: 4
+    mlp_hidden_size: 19456
+    activation_type: swiglu
+    block_type: sequential
+    rope: true
+    rope_full_precision: true
+    rope_theta: 5000000.0
+    rope_type: default
+    rope_factor: null
+    rope_high_freq_factor: null
+    rope_low_freq_factor: null
+    rope_original_max_position_embeddings: null
+    rope_attention_factor: null
+    rope_beta_fast: null
+    rope_beta_slow: null
+    rope_mscale: null
+    rope_mscale_all_dim: null
+    rope_truncate: null
+    attention_type: sdpa
+    full_attention_layers: null
+    sliding_attention_rope_scaling: false
+    float32_attention: true
+    attention_dropout: 0.0
+    attention_layer_norm: true
+    attention_layer_norm_type: qwen3
+    residual_dropout: 0.1
+    response_residual_dropout: 0.0
+    layer_norm_type: rms
+    layer_norm_with_affine: true
+    layer_norm_eps: 1.0e-06
+    attention_layer_norm_with_affine: true
+    max_sequence_length: 16384
+    max_position_embeddings: null
+    include_bias: false
+    bias_for_layer_norm: null
+    norm_after: false
+    moe_num_experts: 8
+    moe_top_k: 2
+    moe_mlp_impl: sparse
+    moe_log_expert_assignment: false
+    moe_shared_expert: false
+    moe_lbl_in_fp32: false
+    moe_interleave: false
+    moe_loss_weight: 0.1
+    moe_zloss_weight: null
+    moe_dropless: true
+    moe_capacity_factor: 1.25
+    embedding_dropout: 0.0
+    scale_logits: false
+    vocab_size: 151669
+    additional_vocab_size: 128
+    weight_tying: false
+    embedding_size: 155648
+    use_position_ids: true
+    tokenizer:
+      identifier: Qwen/Qwen3-4B
+      tokenizer_dir: null
+      new_tokens_for_both_input_and_output:
+      - <setup_start>
+      - <setup_end>
+      - <control_start>
+      - <control_end>
+      - <state_start>
+      - <state_end>
+      - <state_0>
+      - <state_1>
+      - <state_2>
+      - <state_3>
+      - <state_4>
+      - <state_5>
+      - <state_6>
+      - <state_7>
+      - <state_8>
+      - <state_9>
+      - <state_10>
+      - <state_11>
+      - <state_12>
+      - <state_13>
+      - <state_14>
+      - <state_15>
+      - <state_16>
+      - <state_17>
+      - <state_18>
+      - <state_19>
+      - <state_20>
+      - <state_21>
+      - <state_22>
+      - <state_23>
+      - <state_24>
+      - <state_25>
+      - <state_26>
+      - <state_27>
+      - <state_28>
+      - <state_29>
+      - <state_30>
+      - <state_31>
+      - <state_32>
+      - <state_33>
+      - <state_34>
+      - <state_35>
+      - <state_36>
+      - <state_37>
+      - <state_38>
+      - <state_39>
+      - <state_40>
+      - <state_41>
+      - <state_42>
+      - <state_43>
+      - <state_44>
+      - <state_45>
+      - <state_46>
+      - <state_47>
+      - <state_48>
+      - <state_49>
+      - <state_50>
+      - <state_51>
+      - <state_52>
+      - <state_53>
+      - <state_54>
+      - <state_55>
+      - <state_56>
+      - <state_57>
+      - <state_58>
+      - <state_59>
+      - <state_60>
+      - <state_61>
+      - <state_62>
+      - <state_63>
+      - <state_64>
+      - <state_65>
+      - <state_66>
+      - <state_67>
+      - <state_68>
+      - <state_69>
+      - <state_70>
+      - <state_71>
+      - <state_72>
+      - <state_73>
+      - <state_74>
+      - <state_75>
+      - <state_76>
+      - <state_77>
+      - <state_78>
+      - <state_79>
+      - <state_80>
+      - <state_81>
+      - <state_82>
+      - <state_83>
+      - <state_84>
+      - <state_85>
+      - <state_86>
+      - <state_87>
+      - <state_88>
+      - <state_89>
+      - <state_90>
+      - <state_91>
+      - <state_92>
+      - <state_93>
+      - <state_94>
+      - <state_95>
+      - <state_96>
+      - <state_97>
+      - <state_98>
+      - <state_99>
+      - <state_100>
+      - <state_101>
+      - <state_102>
+      - <state_103>
+      - <state_104>
+      - <state_105>
+      - <state_106>
+      - <state_107>
+      - <state_108>
+      - <state_109>
+      - <state_110>
+      - <state_111>
+      - <state_112>
+      - <state_113>
+      - <state_114>
+      - <state_115>
+      - <state_116>
+      - <state_117>
+      - <state_118>
+      - <state_119>
+      - <state_120>
+      - <state_121>
+      - <state_122>
+      - <state_123>
+      - <state_124>
+      - <state_125>
+      - <state_126>
+      - <state_127>
+      - <state_128>
+      - <state_129>
+      - <state_130>
+      - <state_131>
+      - <state_132>
+      - <state_133>
+      - <state_134>
+      - <state_135>
+      - <state_136>
+      - <state_137>
+      - <state_138>
+      - <state_139>
+      - <state_140>
+      - <state_141>
+      - <state_142>
+      - <state_143>
+      - <state_144>
+      - <state_145>
+      - <state_146>
+      - <state_147>
+      - <state_148>
+      - <state_149>
+      - <state_150>
+      - <state_151>
+      - <state_152>
+      - <state_153>
+      - <state_154>
+      - <state_155>
+      - <state_156>
+      - <state_157>
+      - <state_158>
+      - <state_159>
+      - <state_160>
+      - <state_161>
+      - <state_162>
+      - <state_163>
+      - <state_164>
+      - <state_165>
+      - <state_166>
+      - <state_167>
+      - <state_168>
+      - <state_169>
+      - <state_170>
+      - <state_171>
+      - <state_172>
+      - <state_173>
+      - <state_174>
+      - <state_175>
+      - <state_176>
+      - <state_177>
+      - <state_178>
+      - <state_179>
+      - <state_180>
+      - <state_181>
+      - <state_182>
+      - <state_183>
+      - <state_184>
+      - <state_185>
+      - <state_186>
+      - <state_187>
+      - <state_188>
+      - <state_189>
+      - <state_190>
+      - <state_191>
+      - <state_192>
+      - <state_193>
+      - <state_194>
+      - <state_195>
+      - <state_196>
+      - <state_197>
+      - <state_198>
+      - <state_199>
+      - <state_200>
+      - <state_201>
+      - <state_202>
+      - <state_203>
+      - <state_204>
+      - <state_205>
+      - <state_206>
+      - <state_207>
+      - <state_208>
+      - <state_209>
+      - <state_210>
+      - <state_211>
+      - <state_212>
+      - <state_213>
+      - <state_214>
+      - <state_215>
+      - <state_216>
+      - <state_217>
+      - <state_218>
+      - <state_219>
+      - <state_220>
+      - <state_221>
+      - <state_222>
+      - <state_223>
+      - <state_224>
+      - <state_225>
+      - <state_226>
+      - <state_227>
+      - <state_228>
+      - <state_229>
+      - <state_230>
+      - <state_231>
+      - <state_232>
+      - <state_233>
+      - <state_234>
+      - <state_235>
+      - <state_236>
+      - <state_237>
+      - <state_238>
+      - <state_239>
+      - <state_240>
+      - <state_241>
+      - <state_242>
+      - <state_243>
+      - <state_244>
+      - <state_245>
+      - <state_246>
+      - <state_247>
+      - <state_248>
+      - <state_249>
+      - <state_250>
+      - <state_251>
+      - <state_252>
+      - <state_253>
+      - <state_254>
+      - <state_255>
+      - <action_output>
+      - <action_start>
+      - <action_end>
+      - <action_0>
+      - <action_1>
+      - <action_2>
+      - <action_3>
+      - <action_4>
+      - <action_5>
+      - <action_6>
+      - <action_7>
+      - <action_8>
+      - <action_9>
+      - <action_10>
+      - <action_11>
+      - <action_12>
+      - <action_13>
+      - <action_14>
+      - <action_15>
+      - <action_16>
+      - <action_17>
+      - <action_18>
+      - <action_19>
+      - <action_20>
+      - <action_21>
+      - <action_22>
+      - <action_23>
+      - <action_24>
+      - <action_25>
+      - <action_26>
+      - <action_27>
+      - <action_28>
+      - <action_29>
+      - <action_30>
+      - <action_31>
+      - <action_32>
+      - <action_33>
+      - <action_34>
+      - <action_35>
+      - <action_36>
+      - <action_37>
+      - <action_38>
+      - <action_39>
+      - <action_40>
+      - <action_41>
+      - <action_42>
+      - <action_43>
+      - <action_44>
+      - <action_45>
+      - <action_46>
+      - <action_47>
+      - <action_48>
+      - <action_49>
+      - <action_50>
+      - <action_51>
+      - <action_52>
+      - <action_53>
+      - <action_54>
+      - <action_55>
+      - <action_56>
+      - <action_57>
+      - <action_58>
+      - <action_59>
+      - <action_60>
+      - <action_61>
+      - <action_62>
+      - <action_63>
+      - <action_64>
+      - <action_65>
+      - <action_66>
+      - <action_67>
+      - <action_68>
+      - <action_69>
+      - <action_70>
+      - <action_71>
+      - <action_72>
+      - <action_73>
+      - <action_74>
+      - <action_75>
+      - <action_76>
+      - <action_77>
+      - <action_78>
+      - <action_79>
+      - <action_80>
+      - <action_81>
+      - <action_82>
+      - <action_83>
+      - <action_84>
+      - <action_85>
+      - <action_86>
+      - <action_87>
+      - <action_88>
+      - <action_89>
+      - <action_90>
+      - <action_91>
+      - <action_92>
+      - <action_93>
+      - <action_94>
+      - <action_95>
+      - <action_96>
+      - <action_97>
+      - <action_98>
+      - <action_99>
+      - <action_100>
+      - <action_101>
+      - <action_102>
+      - <action_103>
+      - <action_104>
+      - <action_105>
+      - <action_106>
+      - <action_107>
+      - <action_108>
+      - <action_109>
+      - <action_110>
+      - <action_111>
+      - <action_112>
+      - <action_113>
+      - <action_114>
+      - <action_115>
+      - <action_116>
+      - <action_117>
+      - <action_118>
+      - <action_119>
+      - <action_120>
+      - <action_121>
+      - <action_122>
+      - <action_123>
+      - <action_124>
+      - <action_125>
+      - <action_126>
+      - <action_127>
+      - <action_128>
+      - <action_129>
+      - <action_130>
+      - <action_131>
+      - <action_132>
+      - <action_133>
+      - <action_134>
+      - <action_135>
+      - <action_136>
+      - <action_137>
+      - <action_138>
+      - <action_139>
+      - <action_140>
+      - <action_141>
+      - <action_142>
+      - <action_143>
+      - <action_144>
+      - <action_145>
+      - <action_146>
+      - <action_147>
+      - <action_148>
+      - <action_149>
+      - <action_150>
+      - <action_151>
+      - <action_152>
+      - <action_153>
+      - <action_154>
+      - <action_155>
+      - <action_156>
+      - <action_157>
+      - <action_158>
+      - <action_159>
+      - <action_160>
+      - <action_161>
+      - <action_162>
+      - <action_163>
+      - <action_164>
+      - <action_165>
+      - <action_166>
+      - <action_167>
+      - <action_168>
+      - <action_169>
+      - <action_170>
+      - <action_171>
+      - <action_172>
+      - <action_173>
+      - <action_174>
+      - <action_175>
+      - <action_176>
+      - <action_177>
+      - <action_178>
+      - <action_179>
+      - <action_180>
+      - <action_181>
+      - <action_182>
+      - <action_183>
+      - <action_184>
+      - <action_185>
+      - <action_186>
+      - <action_187>
+      - <action_188>
+      - <action_189>
+      - <action_190>
+      - <action_191>
+      - <action_192>
+      - <action_193>
+      - <action_194>
+      - <action_195>
+      - <action_196>
+      - <action_197>
+      - <action_198>
+      - <action_199>
+      - <action_200>
+      - <action_201>
+      - <action_202>
+      - <action_203>
+      - <action_204>
+      - <action_205>
+      - <action_206>
+      - <action_207>
+      - <action_208>
+      - <action_209>
+      - <action_210>
+      - <action_211>
+      - <action_212>
+      - <action_213>
+      - <action_214>
+      - <action_215>
+      - <action_216>
+      - <action_217>
+      - <action_218>
+      - <action_219>
+      - <action_220>
+      - <action_221>
+      - <action_222>
+      - <action_223>
+      - <action_224>
+      - <action_225>
+      - <action_226>
+      - <action_227>
+      - <action_228>
+      - <action_229>
+      - <action_230>
+      - <action_231>
+      - <action_232>
+      - <action_233>
+      - <action_234>
+      - <action_235>
+      - <action_236>
+      - <action_237>
+      - <action_238>
+      - <action_239>
+      - <action_240>
+      - <action_241>
+      - <action_242>
+      - <action_243>
+      - <action_244>
+      - <action_245>
+      - <action_246>
+      - <action_247>
+      - <action_248>
+      - <action_249>
+      - <action_250>
+      - <action_251>
+      - <action_252>
+      - <action_253>
+      - <action_254>
+      - <action_255>
+      - <action_256>
+      - <action_257>
+      - <action_258>
+      - <action_259>
+      - <action_260>
+      - <action_261>
+      - <action_262>
+      - <action_263>
+      - <action_264>
+      - <action_265>
+      - <action_266>
+      - <action_267>
+      - <action_268>
+      - <action_269>
+      - <action_270>
+      - <action_271>
+      - <action_272>
+      - <action_273>
+      - <action_274>
+      - <action_275>
+      - <action_276>
+      - <action_277>
+      - <action_278>
+      - <action_279>
+      - <action_280>
+      - <action_281>
+      - <action_282>
+      - <action_283>
+      - <action_284>
+      - <action_285>
+      - <action_286>
+      - <action_287>
+      - <action_288>
+      - <action_289>
+      - <action_290>
+      - <action_291>
+      - <action_292>
+      - <action_293>
+      - <action_294>
+      - <action_295>
+      - <action_296>
+      - <action_297>
+      - <action_298>
+      - <action_299>
+      - <action_300>
+      - <action_301>
+      - <action_302>
+      - <action_303>
+      - <action_304>
+      - <action_305>
+      - <action_306>
+      - <action_307>
+      - <action_308>
+      - <action_309>
+      - <action_310>
+      - <action_311>
+      - <action_312>
+      - <action_313>
+      - <action_314>
+      - <action_315>
+      - <action_316>
+      - <action_317>
+      - <action_318>
+      - <action_319>
+      - <action_320>
+      - <action_321>
+      - <action_322>
+      - <action_323>
+      - <action_324>
+      - <action_325>
+      - <action_326>
+      - <action_327>
+      - <action_328>
+      - <action_329>
+      - <action_330>
+      - <action_331>
+      - <action_332>
+      - <action_333>
+      - <action_334>
+      - <action_335>
+      - <action_336>
+      - <action_337>
+      - <action_338>
+      - <action_339>
+      - <action_340>
+      - <action_341>
+      - <action_342>
+      - <action_343>
+      - <action_344>
+      - <action_345>
+      - <action_346>
+      - <action_347>
+      - <action_348>
+      - <action_349>
+      - <action_350>
+      - <action_351>
+      - <action_352>
+      - <action_353>
+      - <action_354>
+      - <action_355>
+      - <action_356>
+      - <action_357>
+      - <action_358>
+      - <action_359>
+      - <action_360>
+      - <action_361>
+      - <action_362>
+      - <action_363>
+      - <action_364>
+      - <action_365>
+      - <action_366>
+      - <action_367>
+      - <action_368>
+      - <action_369>
+      - <action_370>
+      - <action_371>
+      - <action_372>
+      - <action_373>
+      - <action_374>
+      - <action_375>
+      - <action_376>
+      - <action_377>
+      - <action_378>
+      - <action_379>
+      - <action_380>
+      - <action_381>
+      - <action_382>
+      - <action_383>
+      - <action_384>
+      - <action_385>
+      - <action_386>
+      - <action_387>
+      - <action_388>
+      - <action_389>
+      - <action_390>
+      - <action_391>
+      - <action_392>
+      - <action_393>
+      - <action_394>
+      - <action_395>
+      - <action_396>
+      - <action_397>
+      - <action_398>
+      - <action_399>
+      - <action_400>
+      - <action_401>
+      - <action_402>
+      - <action_403>
+      - <action_404>
+      - <action_405>
+      - <action_406>
+      - <action_407>
+      - <action_408>
+      - <action_409>
+      - <action_410>
+      - <action_411>
+      - <action_412>
+      - <action_413>
+      - <action_414>
+      - <action_415>
+      - <action_416>
+      - <action_417>
+      - <action_418>
+      - <action_419>
+      - <action_420>
+      - <action_421>
+      - <action_422>
+      - <action_423>
+      - <action_424>
+      - <action_425>
+      - <action_426>
+      - <action_427>
+      - <action_428>
+      - <action_429>
+      - <action_430>
+      - <action_431>
+      - <action_432>
+      - <action_433>
+      - <action_434>
+      - <action_435>
+      - <action_436>
+      - <action_437>
+      - <action_438>
+      - <action_439>
+      - <action_440>
+      - <action_441>
+      - <action_442>
+      - <action_443>
+      - <action_444>
+      - <action_445>
+      - <action_446>
+      - <action_447>
+      - <action_448>
+      - <action_449>
+      - <action_450>
+      - <action_451>
+      - <action_452>
+      - <action_453>
+      - <action_454>
+      - <action_455>
+      - <action_456>
+      - <action_457>
+      - <action_458>
+      - <action_459>
+      - <action_460>
+      - <action_461>
+      - <action_462>
+      - <action_463>
+      - <action_464>
+      - <action_465>
+      - <action_466>
+      - <action_467>
+      - <action_468>
+      - <action_469>
+      - <action_470>
+      - <action_471>
+      - <action_472>
+      - <action_473>
+      - <action_474>
+      - <action_475>
+      - <action_476>
+      - <action_477>
+      - <action_478>
+      - <action_479>
+      - <action_480>
+      - <action_481>
+      - <action_482>
+      - <action_483>
+      - <action_484>
+      - <action_485>
+      - <action_486>
+      - <action_487>
+      - <action_488>
+      - <action_489>
+      - <action_490>
+      - <action_491>
+      - <action_492>
+      - <action_493>
+      - <action_494>
+      - <action_495>
+      - <action_496>
+      - <action_497>
+      - <action_498>
+      - <action_499>
+      - <action_500>
+      - <action_501>
+      - <action_502>
+      - <action_503>
+      - <action_504>
+      - <action_505>
+      - <action_506>
+      - <action_507>
+      - <action_508>
+      - <action_509>
+      - <action_510>
+      - <action_511>
+      - <action_512>
+      - <action_513>
+      - <action_514>
+      - <action_515>
+      - <action_516>
+      - <action_517>
+      - <action_518>
+      - <action_519>
+      - <action_520>
+      - <action_521>
+      - <action_522>
+      - <action_523>
+      - <action_524>
+      - <action_525>
+      - <action_526>
+      - <action_527>
+      - <action_528>
+      - <action_529>
+      - <action_530>
+      - <action_531>
+      - <action_532>
+      - <action_533>
+      - <action_534>
+      - <action_535>
+      - <action_536>
+      - <action_537>
+      - <action_538>
+      - <action_539>
+      - <action_540>
+      - <action_541>
+      - <action_542>
+      - <action_543>
+      - <action_544>
+      - <action_545>
+      - <action_546>
+      - <action_547>
+      - <action_548>
+      - <action_549>
+      - <action_550>
+      - <action_551>
+      - <action_552>
+      - <action_553>
+      - <action_554>
+      - <action_555>
+      - <action_556>
+      - <action_557>
+      - <action_558>
+      - <action_559>
+      - <action_560>
+      - <action_561>
+      - <action_562>
+      - <action_563>
+      - <action_564>
+      - <action_565>
+      - <action_566>
+      - <action_567>
+      - <action_568>
+      - <action_569>
+      - <action_570>
+      - <action_571>
+      - <action_572>
+      - <action_573>
+      - <action_574>
+      - <action_575>
+      - <action_576>
+      - <action_577>
+      - <action_578>
+      - <action_579>
+      - <action_580>
+      - <action_581>
+      - <action_582>
+      - <action_583>
+      - <action_584>
+      - <action_585>
+      - <action_586>
+      - <action_587>
+      - <action_588>
+      - <action_589>
+      - <action_590>
+      - <action_591>
+      - <action_592>
+      - <action_593>
+      - <action_594>
+      - <action_595>
+      - <action_596>
+      - <action_597>
+      - <action_598>
+      - <action_599>
+      - <action_600>
+      - <action_601>
+      - <action_602>
+      - <action_603>
+      - <action_604>
+      - <action_605>
+      - <action_606>
+      - <action_607>
+      - <action_608>
+      - <action_609>
+      - <action_610>
+      - <action_611>
+      - <action_612>
+      - <action_613>
+      - <action_614>
+      - <action_615>
+      - <action_616>
+      - <action_617>
+      - <action_618>
+      - <action_619>
+      - <action_620>
+      - <action_621>
+      - <action_622>
+      - <action_623>
+      - <action_624>
+      - <action_625>
+      - <action_626>
+      - <action_627>
+      - <action_628>
+      - <action_629>
+      - <action_630>
+      - <action_631>
+      - <action_632>
+      - <action_633>
+      - <action_634>
+      - <action_635>
+      - <action_636>
+      - <action_637>
+      - <action_638>
+      - <action_639>
+      - <action_640>
+      - <action_641>
+      - <action_642>
+      - <action_643>
+      - <action_644>
+      - <action_645>
+      - <action_646>
+      - <action_647>
+      - <action_648>
+      - <action_649>
+      - <action_650>
+      - <action_651>
+      - <action_652>
+      - <action_653>
+      - <action_654>
+      - <action_655>
+      - <action_656>
+      - <action_657>
+      - <action_658>
+      - <action_659>
+      - <action_660>
+      - <action_661>
+      - <action_662>
+      - <action_663>
+      - <action_664>
+      - <action_665>
+      - <action_666>
+      - <action_667>
+      - <action_668>
+      - <action_669>
+      - <action_670>
+      - <action_671>
+      - <action_672>
+      - <action_673>
+      - <action_674>
+      - <action_675>
+      - <action_676>
+      - <action_677>
+      - <action_678>
+      - <action_679>
+      - <action_680>
+      - <action_681>
+      - <action_682>
+      - <action_683>
+      - <action_684>
+      - <action_685>
+      - <action_686>
+      - <action_687>
+      - <action_688>
+      - <action_689>
+      - <action_690>
+      - <action_691>
+      - <action_692>
+      - <action_693>
+      - <action_694>
+      - <action_695>
+      - <action_696>
+      - <action_697>
+      - <action_698>
+      - <action_699>
+      - <action_700>
+      - <action_701>
+      - <action_702>
+      - <action_703>
+      - <action_704>
+      - <action_705>
+      - <action_706>
+      - <action_707>
+      - <action_708>
+      - <action_709>
+      - <action_710>
+      - <action_711>
+      - <action_712>
+      - <action_713>
+      - <action_714>
+      - <action_715>
+      - <action_716>
+      - <action_717>
+      - <action_718>
+      - <action_719>
+      - <action_720>
+      - <action_721>
+      - <action_722>
+      - <action_723>
+      - <action_724>
+      - <action_725>
+      - <action_726>
+      - <action_727>
+      - <action_728>
+      - <action_729>
+      - <action_730>
+      - <action_731>
+      - <action_732>
+      - <action_733>
+      - <action_734>
+      - <action_735>
+      - <action_736>
+      - <action_737>
+      - <action_738>
+      - <action_739>
+      - <action_740>
+      - <action_741>
+      - <action_742>
+      - <action_743>
+      - <action_744>
+      - <action_745>
+      - <action_746>
+      - <action_747>
+      - <action_748>
+      - <action_749>
+      - <action_750>
+      - <action_751>
+      - <action_752>
+      - <action_753>
+      - <action_754>
+      - <action_755>
+      - <action_756>
+      - <action_757>
+      - <action_758>
+      - <action_759>
+      - <action_760>
+      - <action_761>
+      - <action_762>
+      - <action_763>
+      - <action_764>
+      - <action_765>
+      - <action_766>
+      - <action_767>
+      - <action_768>
+      - <action_769>
+      - <action_770>
+      - <action_771>
+      - <action_772>
+      - <action_773>
+      - <action_774>
+      - <action_775>
+      - <action_776>
+      - <action_777>
+      - <action_778>
+      - <action_779>
+      - <action_780>
+      - <action_781>
+      - <action_782>
+      - <action_783>
+      - <action_784>
+      - <action_785>
+      - <action_786>
+      - <action_787>
+      - <action_788>
+      - <action_789>
+      - <action_790>
+      - <action_791>
+      - <action_792>
+      - <action_793>
+      - <action_794>
+      - <action_795>
+      - <action_796>
+      - <action_797>
+      - <action_798>
+      - <action_799>
+      - <action_800>
+      - <action_801>
+      - <action_802>
+      - <action_803>
+      - <action_804>
+      - <action_805>
+      - <action_806>
+      - <action_807>
+      - <action_808>
+      - <action_809>
+      - <action_810>
+      - <action_811>
+      - <action_812>
+      - <action_813>
+      - <action_814>
+      - <action_815>
+      - <action_816>
+      - <action_817>
+      - <action_818>
+      - <action_819>
+      - <action_820>
+      - <action_821>
+      - <action_822>
+      - <action_823>
+      - <action_824>
+      - <action_825>
+      - <action_826>
+      - <action_827>
+      - <action_828>
+      - <action_829>
+      - <action_830>
+      - <action_831>
+      - <action_832>
+      - <action_833>
+      - <action_834>
+      - <action_835>
+      - <action_836>
+      - <action_837>
+      - <action_838>
+      - <action_839>
+      - <action_840>
+      - <action_841>
+      - <action_842>
+      - <action_843>
+      - <action_844>
+      - <action_845>
+      - <action_846>
+      - <action_847>
+      - <action_848>
+      - <action_849>
+      - <action_850>
+      - <action_851>
+      - <action_852>
+      - <action_853>
+      - <action_854>
+      - <action_855>
+      - <action_856>
+      - <action_857>
+      - <action_858>
+      - <action_859>
+      - <action_860>
+      - <action_861>
+      - <action_862>
+      - <action_863>
+      - <action_864>
+      - <action_865>
+      - <action_866>
+      - <action_867>
+      - <action_868>
+      - <action_869>
+      - <action_870>
+      - <action_871>
+      - <action_872>
+      - <action_873>
+      - <action_874>
+      - <action_875>
+      - <action_876>
+      - <action_877>
+      - <action_878>
+      - <action_879>
+      - <action_880>
+      - <action_881>
+      - <action_882>
+      - <action_883>
+      - <action_884>
+      - <action_885>
+      - <action_886>
+      - <action_887>
+      - <action_888>
+      - <action_889>
+      - <action_890>
+      - <action_891>
+      - <action_892>
+      - <action_893>
+      - <action_894>
+      - <action_895>
+      - <action_896>
+      - <action_897>
+      - <action_898>
+      - <action_899>
+      - <action_900>
+      - <action_901>
+      - <action_902>
+      - <action_903>
+      - <action_904>
+      - <action_905>
+      - <action_906>
+      - <action_907>
+      - <action_908>
+      - <action_909>
+      - <action_910>
+      - <action_911>
+      - <action_912>
+      - <action_913>
+      - <action_914>
+      - <action_915>
+      - <action_916>
+      - <action_917>
+      - <action_918>
+      - <action_919>
+      - <action_920>
+      - <action_921>
+      - <action_922>
+      - <action_923>
+      - <action_924>
+      - <action_925>
+      - <action_926>
+      - <action_927>
+      - <action_928>
+      - <action_929>
+      - <action_930>
+      - <action_931>
+      - <action_932>
+      - <action_933>
+      - <action_934>
+      - <action_935>
+      - <action_936>
+      - <action_937>
+      - <action_938>
+      - <action_939>
+      - <action_940>
+      - <action_941>
+      - <action_942>
+      - <action_943>
+      - <action_944>
+      - <action_945>
+      - <action_946>
+      - <action_947>
+      - <action_948>
+      - <action_949>
+      - <action_950>
+      - <action_951>
+      - <action_952>
+      - <action_953>
+      - <action_954>
+      - <action_955>
+      - <action_956>
+      - <action_957>
+      - <action_958>
+      - <action_959>
+      - <action_960>
+      - <action_961>
+      - <action_962>
+      - <action_963>
+      - <action_964>
+      - <action_965>
+      - <action_966>
+      - <action_967>
+      - <action_968>
+      - <action_969>
+      - <action_970>
+      - <action_971>
+      - <action_972>
+      - <action_973>
+      - <action_974>
+      - <action_975>
+      - <action_976>
+      - <action_977>
+      - <action_978>
+      - <action_979>
+      - <action_980>
+      - <action_981>
+      - <action_982>
+      - <action_983>
+      - <action_984>
+      - <action_985>
+      - <action_986>
+      - <action_987>
+      - <action_988>
+      - <action_989>
+      - <action_990>
+      - <action_991>
+      - <action_992>
+      - <action_993>
+      - <action_994>
+      - <action_995>
+      - <action_996>
+      - <action_997>
+      - <action_998>
+      - <action_999>
+      - <action_1000>
+      - <action_1001>
+      - <action_1002>
+      - <action_1003>
+      - <action_1004>
+      - <action_1005>
+      - <action_1006>
+      - <action_1007>
+      - <action_1008>
+      - <action_1009>
+      - <action_1010>
+      - <action_1011>
+      - <action_1012>
+      - <action_1013>
+      - <action_1014>
+      - <action_1015>
+      - <action_1016>
+      - <action_1017>
+      - <action_1018>
+      - <action_1019>
+      - <action_1020>
+      - <action_1021>
+      - <action_1022>
+      - <action_1023>
+      - <action_1024>
+      - <action_1025>
+      - <action_1026>
+      - <action_1027>
+      - <action_1028>
+      - <action_1029>
+      - <action_1030>
+      - <action_1031>
+      - <action_1032>
+      - <action_1033>
+      - <action_1034>
+      - <action_1035>
+      - <action_1036>
+      - <action_1037>
+      - <action_1038>
+      - <action_1039>
+      - <action_1040>
+      - <action_1041>
+      - <action_1042>
+      - <action_1043>
+      - <action_1044>
+      - <action_1045>
+      - <action_1046>
+      - <action_1047>
+      - <action_1048>
+      - <action_1049>
+      - <action_1050>
+      - <action_1051>
+      - <action_1052>
+      - <action_1053>
+      - <action_1054>
+      - <action_1055>
+      - <action_1056>
+      - <action_1057>
+      - <action_1058>
+      - <action_1059>
+      - <action_1060>
+      - <action_1061>
+      - <action_1062>
+      - <action_1063>
+      - <action_1064>
+      - <action_1065>
+      - <action_1066>
+      - <action_1067>
+      - <action_1068>
+      - <action_1069>
+      - <action_1070>
+      - <action_1071>
+      - <action_1072>
+      - <action_1073>
+      - <action_1074>
+      - <action_1075>
+      - <action_1076>
+      - <action_1077>
+      - <action_1078>
+      - <action_1079>
+      - <action_1080>
+      - <action_1081>
+      - <action_1082>
+      - <action_1083>
+      - <action_1084>
+      - <action_1085>
+      - <action_1086>
+      - <action_1087>
+      - <action_1088>
+      - <action_1089>
+      - <action_1090>
+      - <action_1091>
+      - <action_1092>
+      - <action_1093>
+      - <action_1094>
+      - <action_1095>
+      - <action_1096>
+      - <action_1097>
+      - <action_1098>
+      - <action_1099>
+      - <action_1100>
+      - <action_1101>
+      - <action_1102>
+      - <action_1103>
+      - <action_1104>
+      - <action_1105>
+      - <action_1106>
+      - <action_1107>
+      - <action_1108>
+      - <action_1109>
+      - <action_1110>
+      - <action_1111>
+      - <action_1112>
+      - <action_1113>
+      - <action_1114>
+      - <action_1115>
+      - <action_1116>
+      - <action_1117>
+      - <action_1118>
+      - <action_1119>
+      - <action_1120>
+      - <action_1121>
+      - <action_1122>
+      - <action_1123>
+      - <action_1124>
+      - <action_1125>
+      - <action_1126>
+      - <action_1127>
+      - <action_1128>
+      - <action_1129>
+      - <action_1130>
+      - <action_1131>
+      - <action_1132>
+      - <action_1133>
+      - <action_1134>
+      - <action_1135>
+      - <action_1136>
+      - <action_1137>
+      - <action_1138>
+      - <action_1139>
+      - <action_1140>
+      - <action_1141>
+      - <action_1142>
+      - <action_1143>
+      - <action_1144>
+      - <action_1145>
+      - <action_1146>
+      - <action_1147>
+      - <action_1148>
+      - <action_1149>
+      - <action_1150>
+      - <action_1151>
+      - <action_1152>
+      - <action_1153>
+      - <action_1154>
+      - <action_1155>
+      - <action_1156>
+      - <action_1157>
+      - <action_1158>
+      - <action_1159>
+      - <action_1160>
+      - <action_1161>
+      - <action_1162>
+      - <action_1163>
+      - <action_1164>
+      - <action_1165>
+      - <action_1166>
+      - <action_1167>
+      - <action_1168>
+      - <action_1169>
+      - <action_1170>
+      - <action_1171>
+      - <action_1172>
+      - <action_1173>
+      - <action_1174>
+      - <action_1175>
+      - <action_1176>
+      - <action_1177>
+      - <action_1178>
+      - <action_1179>
+      - <action_1180>
+      - <action_1181>
+      - <action_1182>
+      - <action_1183>
+      - <action_1184>
+      - <action_1185>
+      - <action_1186>
+      - <action_1187>
+      - <action_1188>
+      - <action_1189>
+      - <action_1190>
+      - <action_1191>
+      - <action_1192>
+      - <action_1193>
+      - <action_1194>
+      - <action_1195>
+      - <action_1196>
+      - <action_1197>
+      - <action_1198>
+      - <action_1199>
+      - <action_1200>
+      - <action_1201>
+      - <action_1202>
+      - <action_1203>
+      - <action_1204>
+      - <action_1205>
+      - <action_1206>
+      - <action_1207>
+      - <action_1208>
+      - <action_1209>
+      - <action_1210>
+      - <action_1211>
+      - <action_1212>
+      - <action_1213>
+      - <action_1214>
+      - <action_1215>
+      - <action_1216>
+      - <action_1217>
+      - <action_1218>
+      - <action_1219>
+      - <action_1220>
+      - <action_1221>
+      - <action_1222>
+      - <action_1223>
+      - <action_1224>
+      - <action_1225>
+      - <action_1226>
+      - <action_1227>
+      - <action_1228>
+      - <action_1229>
+      - <action_1230>
+      - <action_1231>
+      - <action_1232>
+      - <action_1233>
+      - <action_1234>
+      - <action_1235>
+      - <action_1236>
+      - <action_1237>
+      - <action_1238>
+      - <action_1239>
+      - <action_1240>
+      - <action_1241>
+      - <action_1242>
+      - <action_1243>
+      - <action_1244>
+      - <action_1245>
+      - <action_1246>
+      - <action_1247>
+      - <action_1248>
+      - <action_1249>
+      - <action_1250>
+      - <action_1251>
+      - <action_1252>
+      - <action_1253>
+      - <action_1254>
+      - <action_1255>
+      - <action_1256>
+      - <action_1257>
+      - <action_1258>
+      - <action_1259>
+      - <action_1260>
+      - <action_1261>
+      - <action_1262>
+      - <action_1263>
+      - <action_1264>
+      - <action_1265>
+      - <action_1266>
+      - <action_1267>
+      - <action_1268>
+      - <action_1269>
+      - <action_1270>
+      - <action_1271>
+      - <action_1272>
+      - <action_1273>
+      - <action_1274>
+      - <action_1275>
+      - <action_1276>
+      - <action_1277>
+      - <action_1278>
+      - <action_1279>
+      - <action_1280>
+      - <action_1281>
+      - <action_1282>
+      - <action_1283>
+      - <action_1284>
+      - <action_1285>
+      - <action_1286>
+      - <action_1287>
+      - <action_1288>
+      - <action_1289>
+      - <action_1290>
+      - <action_1291>
+      - <action_1292>
+      - <action_1293>
+      - <action_1294>
+      - <action_1295>
+      - <action_1296>
+      - <action_1297>
+      - <action_1298>
+      - <action_1299>
+      - <action_1300>
+      - <action_1301>
+      - <action_1302>
+      - <action_1303>
+      - <action_1304>
+      - <action_1305>
+      - <action_1306>
+      - <action_1307>
+      - <action_1308>
+      - <action_1309>
+      - <action_1310>
+      - <action_1311>
+      - <action_1312>
+      - <action_1313>
+      - <action_1314>
+      - <action_1315>
+      - <action_1316>
+      - <action_1317>
+      - <action_1318>
+      - <action_1319>
+      - <action_1320>
+      - <action_1321>
+      - <action_1322>
+      - <action_1323>
+      - <action_1324>
+      - <action_1325>
+      - <action_1326>
+      - <action_1327>
+      - <action_1328>
+      - <action_1329>
+      - <action_1330>
+      - <action_1331>
+      - <action_1332>
+      - <action_1333>
+      - <action_1334>
+      - <action_1335>
+      - <action_1336>
+      - <action_1337>
+      - <action_1338>
+      - <action_1339>
+      - <action_1340>
+      - <action_1341>
+      - <action_1342>
+      - <action_1343>
+      - <action_1344>
+      - <action_1345>
+      - <action_1346>
+      - <action_1347>
+      - <action_1348>
+      - <action_1349>
+      - <action_1350>
+      - <action_1351>
+      - <action_1352>
+      - <action_1353>
+      - <action_1354>
+      - <action_1355>
+      - <action_1356>
+      - <action_1357>
+      - <action_1358>
+      - <action_1359>
+      - <action_1360>
+      - <action_1361>
+      - <action_1362>
+      - <action_1363>
+      - <action_1364>
+      - <action_1365>
+      - <action_1366>
+      - <action_1367>
+      - <action_1368>
+      - <action_1369>
+      - <action_1370>
+      - <action_1371>
+      - <action_1372>
+      - <action_1373>
+      - <action_1374>
+      - <action_1375>
+      - <action_1376>
+      - <action_1377>
+      - <action_1378>
+      - <action_1379>
+      - <action_1380>
+      - <action_1381>
+      - <action_1382>
+      - <action_1383>
+      - <action_1384>
+      - <action_1385>
+      - <action_1386>
+      - <action_1387>
+      - <action_1388>
+      - <action_1389>
+      - <action_1390>
+      - <action_1391>
+      - <action_1392>
+      - <action_1393>
+      - <action_1394>
+      - <action_1395>
+      - <action_1396>
+      - <action_1397>
+      - <action_1398>
+      - <action_1399>
+      - <action_1400>
+      - <action_1401>
+      - <action_1402>
+      - <action_1403>
+      - <action_1404>
+      - <action_1405>
+      - <action_1406>
+      - <action_1407>
+      - <action_1408>
+      - <action_1409>
+      - <action_1410>
+      - <action_1411>
+      - <action_1412>
+      - <action_1413>
+      - <action_1414>
+      - <action_1415>
+      - <action_1416>
+      - <action_1417>
+      - <action_1418>
+      - <action_1419>
+      - <action_1420>
+      - <action_1421>
+      - <action_1422>
+      - <action_1423>
+      - <action_1424>
+      - <action_1425>
+      - <action_1426>
+      - <action_1427>
+      - <action_1428>
+      - <action_1429>
+      - <action_1430>
+      - <action_1431>
+      - <action_1432>
+      - <action_1433>
+      - <action_1434>
+      - <action_1435>
+      - <action_1436>
+      - <action_1437>
+      - <action_1438>
+      - <action_1439>
+      - <action_1440>
+      - <action_1441>
+      - <action_1442>
+      - <action_1443>
+      - <action_1444>
+      - <action_1445>
+      - <action_1446>
+      - <action_1447>
+      - <action_1448>
+      - <action_1449>
+      - <action_1450>
+      - <action_1451>
+      - <action_1452>
+      - <action_1453>
+      - <action_1454>
+      - <action_1455>
+      - <action_1456>
+      - <action_1457>
+      - <action_1458>
+      - <action_1459>
+      - <action_1460>
+      - <action_1461>
+      - <action_1462>
+      - <action_1463>
+      - <action_1464>
+      - <action_1465>
+      - <action_1466>
+      - <action_1467>
+      - <action_1468>
+      - <action_1469>
+      - <action_1470>
+      - <action_1471>
+      - <action_1472>
+      - <action_1473>
+      - <action_1474>
+      - <action_1475>
+      - <action_1476>
+      - <action_1477>
+      - <action_1478>
+      - <action_1479>
+      - <action_1480>
+      - <action_1481>
+      - <action_1482>
+      - <action_1483>
+      - <action_1484>
+      - <action_1485>
+      - <action_1486>
+      - <action_1487>
+      - <action_1488>
+      - <action_1489>
+      - <action_1490>
+      - <action_1491>
+      - <action_1492>
+      - <action_1493>
+      - <action_1494>
+      - <action_1495>
+      - <action_1496>
+      - <action_1497>
+      - <action_1498>
+      - <action_1499>
+      - <action_1500>
+      - <action_1501>
+      - <action_1502>
+      - <action_1503>
+      - <action_1504>
+      - <action_1505>
+      - <action_1506>
+      - <action_1507>
+      - <action_1508>
+      - <action_1509>
+      - <action_1510>
+      - <action_1511>
+      - <action_1512>
+      - <action_1513>
+      - <action_1514>
+      - <action_1515>
+      - <action_1516>
+      - <action_1517>
+      - <action_1518>
+      - <action_1519>
+      - <action_1520>
+      - <action_1521>
+      - <action_1522>
+      - <action_1523>
+      - <action_1524>
+      - <action_1525>
+      - <action_1526>
+      - <action_1527>
+      - <action_1528>
+      - <action_1529>
+      - <action_1530>
+      - <action_1531>
+      - <action_1532>
+      - <action_1533>
+      - <action_1534>
+      - <action_1535>
+      - <action_1536>
+      - <action_1537>
+      - <action_1538>
+      - <action_1539>
+      - <action_1540>
+      - <action_1541>
+      - <action_1542>
+      - <action_1543>
+      - <action_1544>
+      - <action_1545>
+      - <action_1546>
+      - <action_1547>
+      - <action_1548>
+      - <action_1549>
+      - <action_1550>
+      - <action_1551>
+      - <action_1552>
+      - <action_1553>
+      - <action_1554>
+      - <action_1555>
+      - <action_1556>
+      - <action_1557>
+      - <action_1558>
+      - <action_1559>
+      - <action_1560>
+      - <action_1561>
+      - <action_1562>
+      - <action_1563>
+      - <action_1564>
+      - <action_1565>
+      - <action_1566>
+      - <action_1567>
+      - <action_1568>
+      - <action_1569>
+      - <action_1570>
+      - <action_1571>
+      - <action_1572>
+      - <action_1573>
+      - <action_1574>
+      - <action_1575>
+      - <action_1576>
+      - <action_1577>
+      - <action_1578>
+      - <action_1579>
+      - <action_1580>
+      - <action_1581>
+      - <action_1582>
+      - <action_1583>
+      - <action_1584>
+      - <action_1585>
+      - <action_1586>
+      - <action_1587>
+      - <action_1588>
+      - <action_1589>
+      - <action_1590>
+      - <action_1591>
+      - <action_1592>
+      - <action_1593>
+      - <action_1594>
+      - <action_1595>
+      - <action_1596>
+      - <action_1597>
+      - <action_1598>
+      - <action_1599>
+      - <action_1600>
+      - <action_1601>
+      - <action_1602>
+      - <action_1603>
+      - <action_1604>
+      - <action_1605>
+      - <action_1606>
+      - <action_1607>
+      - <action_1608>
+      - <action_1609>
+      - <action_1610>
+      - <action_1611>
+      - <action_1612>
+      - <action_1613>
+      - <action_1614>
+      - <action_1615>
+      - <action_1616>
+      - <action_1617>
+      - <action_1618>
+      - <action_1619>
+      - <action_1620>
+      - <action_1621>
+      - <action_1622>
+      - <action_1623>
+      - <action_1624>
+      - <action_1625>
+      - <action_1626>
+      - <action_1627>
+      - <action_1628>
+      - <action_1629>
+      - <action_1630>
+      - <action_1631>
+      - <action_1632>
+      - <action_1633>
+      - <action_1634>
+      - <action_1635>
+      - <action_1636>
+      - <action_1637>
+      - <action_1638>
+      - <action_1639>
+      - <action_1640>
+      - <action_1641>
+      - <action_1642>
+      - <action_1643>
+      - <action_1644>
+      - <action_1645>
+      - <action_1646>
+      - <action_1647>
+      - <action_1648>
+      - <action_1649>
+      - <action_1650>
+      - <action_1651>
+      - <action_1652>
+      - <action_1653>
+      - <action_1654>
+      - <action_1655>
+      - <action_1656>
+      - <action_1657>
+      - <action_1658>
+      - <action_1659>
+      - <action_1660>
+      - <action_1661>
+      - <action_1662>
+      - <action_1663>
+      - <action_1664>
+      - <action_1665>
+      - <action_1666>
+      - <action_1667>
+      - <action_1668>
+      - <action_1669>
+      - <action_1670>
+      - <action_1671>
+      - <action_1672>
+      - <action_1673>
+      - <action_1674>
+      - <action_1675>
+      - <action_1676>
+      - <action_1677>
+      - <action_1678>
+      - <action_1679>
+      - <action_1680>
+      - <action_1681>
+      - <action_1682>
+      - <action_1683>
+      - <action_1684>
+      - <action_1685>
+      - <action_1686>
+      - <action_1687>
+      - <action_1688>
+      - <action_1689>
+      - <action_1690>
+      - <action_1691>
+      - <action_1692>
+      - <action_1693>
+      - <action_1694>
+      - <action_1695>
+      - <action_1696>
+      - <action_1697>
+      - <action_1698>
+      - <action_1699>
+      - <action_1700>
+      - <action_1701>
+      - <action_1702>
+      - <action_1703>
+      - <action_1704>
+      - <action_1705>
+      - <action_1706>
+      - <action_1707>
+      - <action_1708>
+      - <action_1709>
+      - <action_1710>
+      - <action_1711>
+      - <action_1712>
+      - <action_1713>
+      - <action_1714>
+      - <action_1715>
+      - <action_1716>
+      - <action_1717>
+      - <action_1718>
+      - <action_1719>
+      - <action_1720>
+      - <action_1721>
+      - <action_1722>
+      - <action_1723>
+      - <action_1724>
+      - <action_1725>
+      - <action_1726>
+      - <action_1727>
+      - <action_1728>
+      - <action_1729>
+      - <action_1730>
+      - <action_1731>
+      - <action_1732>
+      - <action_1733>
+      - <action_1734>
+      - <action_1735>
+      - <action_1736>
+      - <action_1737>
+      - <action_1738>
+      - <action_1739>
+      - <action_1740>
+      - <action_1741>
+      - <action_1742>
+      - <action_1743>
+      - <action_1744>
+      - <action_1745>
+      - <action_1746>
+      - <action_1747>
+      - <action_1748>
+      - <action_1749>
+      - <action_1750>
+      - <action_1751>
+      - <action_1752>
+      - <action_1753>
+      - <action_1754>
+      - <action_1755>
+      - <action_1756>
+      - <action_1757>
+      - <action_1758>
+      - <action_1759>
+      - <action_1760>
+      - <action_1761>
+      - <action_1762>
+      - <action_1763>
+      - <action_1764>
+      - <action_1765>
+      - <action_1766>
+      - <action_1767>
+      - <action_1768>
+      - <action_1769>
+      - <action_1770>
+      - <action_1771>
+      - <action_1772>
+      - <action_1773>
+      - <action_1774>
+      - <action_1775>
+      - <action_1776>
+      - <action_1777>
+      - <action_1778>
+      - <action_1779>
+      - <action_1780>
+      - <action_1781>
+      - <action_1782>
+      - <action_1783>
+      - <action_1784>
+      - <action_1785>
+      - <action_1786>
+      - <action_1787>
+      - <action_1788>
+      - <action_1789>
+      - <action_1790>
+      - <action_1791>
+      - <action_1792>
+      - <action_1793>
+      - <action_1794>
+      - <action_1795>
+      - <action_1796>
+      - <action_1797>
+      - <action_1798>
+      - <action_1799>
+      - <action_1800>
+      - <action_1801>
+      - <action_1802>
+      - <action_1803>
+      - <action_1804>
+      - <action_1805>
+      - <action_1806>
+      - <action_1807>
+      - <action_1808>
+      - <action_1809>
+      - <action_1810>
+      - <action_1811>
+      - <action_1812>
+      - <action_1813>
+      - <action_1814>
+      - <action_1815>
+      - <action_1816>
+      - <action_1817>
+      - <action_1818>
+      - <action_1819>
+      - <action_1820>
+      - <action_1821>
+      - <action_1822>
+      - <action_1823>
+      - <action_1824>
+      - <action_1825>
+      - <action_1826>
+      - <action_1827>
+      - <action_1828>
+      - <action_1829>
+      - <action_1830>
+      - <action_1831>
+      - <action_1832>
+      - <action_1833>
+      - <action_1834>
+      - <action_1835>
+      - <action_1836>
+      - <action_1837>
+      - <action_1838>
+      - <action_1839>
+      - <action_1840>
+      - <action_1841>
+      - <action_1842>
+      - <action_1843>
+      - <action_1844>
+      - <action_1845>
+      - <action_1846>
+      - <action_1847>
+      - <action_1848>
+      - <action_1849>
+      - <action_1850>
+      - <action_1851>
+      - <action_1852>
+      - <action_1853>
+      - <action_1854>
+      - <action_1855>
+      - <action_1856>
+      - <action_1857>
+      - <action_1858>
+      - <action_1859>
+      - <action_1860>
+      - <action_1861>
+      - <action_1862>
+      - <action_1863>
+      - <action_1864>
+      - <action_1865>
+      - <action_1866>
+      - <action_1867>
+      - <action_1868>
+      - <action_1869>
+      - <action_1870>
+      - <action_1871>
+      - <action_1872>
+      - <action_1873>
+      - <action_1874>
+      - <action_1875>
+      - <action_1876>
+      - <action_1877>
+      - <action_1878>
+      - <action_1879>
+      - <action_1880>
+      - <action_1881>
+      - <action_1882>
+      - <action_1883>
+      - <action_1884>
+      - <action_1885>
+      - <action_1886>
+      - <action_1887>
+      - <action_1888>
+      - <action_1889>
+      - <action_1890>
+      - <action_1891>
+      - <action_1892>
+      - <action_1893>
+      - <action_1894>
+      - <action_1895>
+      - <action_1896>
+      - <action_1897>
+      - <action_1898>
+      - <action_1899>
+      - <action_1900>
+      - <action_1901>
+      - <action_1902>
+      - <action_1903>
+      - <action_1904>
+      - <action_1905>
+      - <action_1906>
+      - <action_1907>
+      - <action_1908>
+      - <action_1909>
+      - <action_1910>
+      - <action_1911>
+      - <action_1912>
+      - <action_1913>
+      - <action_1914>
+      - <action_1915>
+      - <action_1916>
+      - <action_1917>
+      - <action_1918>
+      - <action_1919>
+      - <action_1920>
+      - <action_1921>
+      - <action_1922>
+      - <action_1923>
+      - <action_1924>
+      - <action_1925>
+      - <action_1926>
+      - <action_1927>
+      - <action_1928>
+      - <action_1929>
+      - <action_1930>
+      - <action_1931>
+      - <action_1932>
+      - <action_1933>
+      - <action_1934>
+      - <action_1935>
+      - <action_1936>
+      - <action_1937>
+      - <action_1938>
+      - <action_1939>
+      - <action_1940>
+      - <action_1941>
+      - <action_1942>
+      - <action_1943>
+      - <action_1944>
+      - <action_1945>
+      - <action_1946>
+      - <action_1947>
+      - <action_1948>
+      - <action_1949>
+      - <action_1950>
+      - <action_1951>
+      - <action_1952>
+      - <action_1953>
+      - <action_1954>
+      - <action_1955>
+      - <action_1956>
+      - <action_1957>
+      - <action_1958>
+      - <action_1959>
+      - <action_1960>
+      - <action_1961>
+      - <action_1962>
+      - <action_1963>
+      - <action_1964>
+      - <action_1965>
+      - <action_1966>
+      - <action_1967>
+      - <action_1968>
+      - <action_1969>
+      - <action_1970>
+      - <action_1971>
+      - <action_1972>
+      - <action_1973>
+      - <action_1974>
+      - <action_1975>
+      - <action_1976>
+      - <action_1977>
+      - <action_1978>
+      - <action_1979>
+      - <action_1980>
+      - <action_1981>
+      - <action_1982>
+      - <action_1983>
+      - <action_1984>
+      - <action_1985>
+      - <action_1986>
+      - <action_1987>
+      - <action_1988>
+      - <action_1989>
+      - <action_1990>
+      - <action_1991>
+      - <action_1992>
+      - <action_1993>
+      - <action_1994>
+      - <action_1995>
+      - <action_1996>
+      - <action_1997>
+      - <action_1998>
+      - <action_1999>
+      - <action_2000>
+      - <action_2001>
+      - <action_2002>
+      - <action_2003>
+      - <action_2004>
+      - <action_2005>
+      - <action_2006>
+      - <action_2007>
+      - <action_2008>
+      - <action_2009>
+      - <action_2010>
+      - <action_2011>
+      - <action_2012>
+      - <action_2013>
+      - <action_2014>
+      - <action_2015>
+      - <action_2016>
+      - <action_2017>
+      - <action_2018>
+      - <action_2019>
+      - <action_2020>
+      - <action_2021>
+      - <action_2022>
+      - <action_2023>
+      - <action_2024>
+      - <action_2025>
+      - <action_2026>
+      - <action_2027>
+      - <action_2028>
+      - <action_2029>
+      - <action_2030>
+      - <action_2031>
+      - <action_2032>
+      - <action_2033>
+      - <action_2034>
+      - <action_2035>
+      - <action_2036>
+      - <action_2037>
+      - <action_2038>
+      - <action_2039>
+      - <action_2040>
+      - <action_2041>
+      - <action_2042>
+      - <action_2043>
+      - <action_2044>
+      - <action_2045>
+      - <action_2046>
+      - <action_2047>
+      - <depth_output>
+      - <depth_start>
+      - <depth_end>
+      - <depth_0>
+      - <depth_1>
+      - <depth_2>
+      - <depth_3>
+      - <depth_4>
+      - <depth_5>
+      - <depth_6>
+      - <depth_7>
+      - <depth_8>
+      - <depth_9>
+      - <depth_10>
+      - <depth_11>
+      - <depth_12>
+      - <depth_13>
+      - <depth_14>
+      - <depth_15>
+      - <depth_16>
+      - <depth_17>
+      - <depth_18>
+      - <depth_19>
+      - <depth_20>
+      - <depth_21>
+      - <depth_22>
+      - <depth_23>
+      - <depth_24>
+      - <depth_25>
+      - <depth_26>
+      - <depth_27>
+      - <depth_28>
+      - <depth_29>
+      - <depth_30>
+      - <depth_31>
+      - <depth_32>
+      - <depth_33>
+      - <depth_34>
+      - <depth_35>
+      - <depth_36>
+      - <depth_37>
+      - <depth_38>
+      - <depth_39>
+      - <depth_40>
+      - <depth_41>
+      - <depth_42>
+      - <depth_43>
+      - <depth_44>
+      - <depth_45>
+      - <depth_46>
+      - <depth_47>
+      - <depth_48>
+      - <depth_49>
+      - <depth_50>
+      - <depth_51>
+      - <depth_52>
+      - <depth_53>
+      - <depth_54>
+      - <depth_55>
+      - <depth_56>
+      - <depth_57>
+      - <depth_58>
+      - <depth_59>
+      - <depth_60>
+      - <depth_61>
+      - <depth_62>
+      - <depth_63>
+      - <depth_64>
+      - <depth_65>
+      - <depth_66>
+      - <depth_67>
+      - <depth_68>
+      - <depth_69>
+      - <depth_70>
+      - <depth_71>
+      - <depth_72>
+      - <depth_73>
+      - <depth_74>
+      - <depth_75>
+      - <depth_76>
+      - <depth_77>
+      - <depth_78>
+      - <depth_79>
+      - <depth_80>
+      - <depth_81>
+      - <depth_82>
+      - <depth_83>
+      - <depth_84>
+      - <depth_85>
+      - <depth_86>
+      - <depth_87>
+      - <depth_88>
+      - <depth_89>
+      - <depth_90>
+      - <depth_91>
+      - <depth_92>
+      - <depth_93>
+      - <depth_94>
+      - <depth_95>
+      - <depth_96>
+      - <depth_97>
+      - <depth_98>
+      - <depth_99>
+      - <depth_100>
+      - <depth_101>
+      - <depth_102>
+      - <depth_103>
+      - <depth_104>
+      - <depth_105>
+      - <depth_106>
+      - <depth_107>
+      - <depth_108>
+      - <depth_109>
+      - <depth_110>
+      - <depth_111>
+      - <depth_112>
+      - <depth_113>
+      - <depth_114>
+      - <depth_115>
+      - <depth_116>
+      - <depth_117>
+      - <depth_118>
+      - <depth_119>
+      - <depth_120>
+      - <depth_121>
+      - <depth_122>
+      - <depth_123>
+      - <depth_124>
+      - <depth_125>
+      - <depth_126>
+      - <depth_127>
+    init_path: null
+    init_incremental: null
+    new_embedding_init_range: 0.02
+    initializer_range: 0.02
+    normalize_input_embeds: false
+    activation_checkpoint: whole_layer
+    compile: blocks
+    fix_pad_tokenizer: true
+    init_std: 0.02
+    init_fn: normal
+    init_cutoff_factor: null
+  vision_backbone:
+    vit:
+      image_model_type: siglip
+      image_default_input_size:
+      - 378
+      - 378
+      image_patch_size: 14
+      image_pos_patch_size: 14
+      image_emb_dim: 1152
+      image_num_heads: 16
+      image_num_key_value_heads: 16
+      image_num_layers: 27
+      image_head_dim: 72
+      image_mlp_dim: 4304
+      image_mlp_activations: gelu_pytorch_tanh
+      image_dropout_rate: 0.0
+      image_num_pos: 729
+      image_norm_eps: 1.0e-06
+      attention_dropout: 0.0
+      residual_dropout: 0.0
+      initializer_range: 0.02
+      float32_attention: true
+      attention_type: sdpa
+      sdpa_backend: all
+      activation_checkpointing: true
+      init_path: null
+      resize_mode: siglip
+      pad_value: 0.0
+      normalize: siglip
+    image_pooling_2d: attention_meanq
+    pooling_attention_mask: true
+    image_projector: mlp
+    image_padding_embed: null
+    vit_layers:
+    - -3
+    - -9
+    skip_unused_layers: true
+    use_deepstack: false
+    share_connector: false
+    image_feature_dropout: 0.0
+    connector_activation_checkpointing: true
+    compile_vit: blocks
+    pool_size_embeds: null
+    compile_connector: dynamic
+    normalize_on_gpu: true
+    use_image_augmentation: full
+  mm_preprocessor:
+    max_answer_len: null
+    last_message_loss_only: false
+    max_text_tokens: null
+    loss_token_weighting: root_subsegments_root_tokens
+    max_frames: 8
+    frame_sample_mode: uniform_last_frame
+    candidate_sampling_fps:
+    - 0.25
+    - 0.5
+    - 1.0
+    - 2.0
+    - 4.0
+    - 6.0
+    - 8.0
+    - 16.0
+    cache_videos: true
+    loading_method: av
+    max_fps:
+    - 2.0
+    time_sampling: true
+    time_mode: per-frame-compact
+    subtitle_mode: frame_1
+    max_crops: 1
+    overlap_margins:
+    - 4.0
+    - 4.0
+    use_col_tokens: false
+    periodic_high_res_frame: null
+    high_low_train_mode: local_rnd
+    high_res_frame_sample_options: null
+    periodic_sample_rate_training:
+      4:
+      - 0.9
+      - 0.03
+      - 0.03
+      - 0.04
+      3:
+      - 0.6
+      - 0.2
+      - 0.2
+    skip_low_res_in_high_low: false
+    pooling_w: 3
+    pooling_h: 3
+    high_res_pooling_w: null
+    high_res_pooling_h: null
+    query_based_resolution_selection: false
+    max_queries_for_resolution_selection: 8
+    use_frame_special_tokens: true
+    frame_sel_clip_identifier: google/siglip2-so400m-patch14-384
+    image_padding_mask: false
+    max_subtitle_tokens: null
+    image:
+      crop_mode: resize
+      use_col_tokens: true
+      max_crops: 8
+      high_res_max_crops: 24
+      p_high_res: 0.0
+      pooling_w: 2
+      pooling_h: 2
+      overlap_margins:
+      - 4
+      - 4
+      max_images: 5
+      max_multi_image_crops: 8
+      multi_image_pooling_w: 2
+      multi_image_pooling_h: 2
+      use_single_crop_col_tokens: false
+      use_single_crop_start_token: true
+    topk: null
+    prune_from_frame: 0
+  bi_directional_attn: null
+  shared_low_high_embedding: true
+  debug: null
+  cp_enabled: false
+  apply_cp_to_vision_backbone: false
+  lora_enable: false
+  lora_rank: 64
+  lora_alpha: 16
+  lora_dropout: 0.05
+  lora_bias: none
+  max_action_dim: 16
+  action_horizon: 50
+  n_action_steps: null
+  n_obs_steps: 1
+  action_expert:
+    max_horizon: 50
+    max_action_dim: 16
+    hidden_size: 768
+    num_layers: 36
+    num_heads: 8
+    mlp_ratio: 4.0
+    ffn_multiple_of: 256
+    timestep_embed_dim: 256
+    dropout: 0.0
+    attn_dropout: 0.0
+    context_layer_norm: true
+    qk_norm: true
+    qk_norm_eps: 1.0e-06
+    rope: true
+    rope_on_cross_attention: false
+    causal_attn: false
+    memory_enabled: true
+    compile: null
+  memory:
+    enabled: true
+    mode: finetune
+    num_latents: 256
+    memory_dim: 1024
+    cosmos_latent_dim: 48
+    recent_dim: 768
+    action_dim: 16
+    proprio_dim: 16
+    language_dim: 2048
+    tokenizer_hidden_dim: 512
+    num_cross_blocks: 2
+    num_heads: 8
+    stage1_checkpoint: artifacts/checkpoints/rmbench-full-attention/stage1/stage1_step160000.npz
+  add_action_expert: true
+  action_expert_detach_vlm: false
+  action_expert_depth_gate: false
+  action_expert_depth_gate_per_layer: false
+  action_expert_depth_gate_init_bias: -4.0
+  action_format: continuous
+  state_format: discrete
+  flow_matching_num_steps: 10
+  flow_matching_cutoff: 1.0
+  flow_matching_time_offset: 0.001
+  flow_matching_time_scale: 0.999
+  flow_matching_beta_alpha: 1.0
+  flow_matching_beta_beta: 1.5
+  num_flow_timesteps: 1
+  mask_action_chunk_padding: true
+  mask_action_dim_padding: true
+  enable_depth_reasoning: false
+  num_depth_codes: 100
+  depth_code_input_noise_rate: 0.0
+  robot_processor:
+    metadata_by_tag:
+      rmbench_stage2:
+        action_key: action
+        state_keys:
+        - observation.state
+        camera_keys:
+        - observation.images.image
+        - observation.images.wrist_image
+        - observation.images.right_image
+        normalize_gripper: true
+        action_dim: 16
+        action_horizon: 50
+        n_action_steps: 50
+        setup_type: three-camera dual-arm robot in RMBench
+        control_mode: absolute joint pose
+        action_stats:
+          mean:
+          - -0.2884466350078583
+          - 1.116854190826416
+          - 0.8176121115684509
+          - -0.6286301016807556
+          - 0.0027539930306375027
+          - 0.3858006000518799
+          - 0.0
+          - 0.29012763500213623
+          - 1.1434277296066284
+          - 0.834911048412323
+          - -0.6164064407348633
+          - 0.0033871293999254704
+          - 0.1835649162530899
+          - 0.0
+          - 0.44485998153686523
+          - 0.693847119808197
+          q01:
+          - -1.1566282510757446
+          - -1.813583730836399e-05
+          - -3.5173223295714706e-05
+          - -1.493809461593628
+          - -0.044773634523153305
+          - -2.2744300365448
+          - 0.0
+          - -2.3844564566388726e-05
+          - -1.9354150936123915e-05
+          - -3.23197637044359e-05
+          - -1.4539425373077393
+          - -0.11438962072134018
+          - -1.3077747821807861
+          - 0.0
+          - 0.0
+          - 0.0
+          q99:
+          - 2.0988931282772683e-05
+          - 2.410673141479492
+          - 2.2643818855285645
+          - 1.7526803276268765e-05
+          - 0.09094329178333282
+          - 1.4207451343536377
+          - 0.0
+          - 0.9992868304252625
+          - 2.5619561672210693
+          - 2.3963286876678467
+          - 2.7463351216283627e-05
+          - 0.12538659572601318
+          - 1.50703763961792
+          - 0.0
+          - 1.0
+          - 1.0
+          std:
+          - 0.2924744486808777
+          - 0.9123055338859558
+          - 0.6921584606170654
+          - 0.5167360901832581
+          - 0.024822130799293518
+          - 0.7730752229690552
+          - 0.0
+          - 0.3022353947162628
+          - 0.9363216161727905
+          - 0.7197534441947937
+          - 0.5159713625907898
+          - 0.042254358530044556
+          - 0.4932859241962433
+          - 0.0
+          - 0.4852803945541382
+          - 0.43509986996650696
+          count:
+          - 277350.0
+          names:
+          - '0'
+          - '1'
+          - '2'
+          - '3'
+          - '4'
+          - '5'
+          - '6'
+          - '7'
+          - '8'
+          - '9'
+          - '10'
+          - '11'
+          - '12'
+          - '13'
+          - '14'
+          - '15'
+          mask:
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+        state_stats:
+          mean:
+          - -0.28619107604026794
+          - 1.1170974969863892
+          - 0.8187518119812012
+          - -0.6292704939842224
+          - 0.002742639509961009
+          - 0.38321176171302795
+          - 0.0
+          - 0.287190318107605
+          - 1.1249585151672363
+          - 0.8228866457939148
+          - -0.6072733998298645
+          - 0.003468270180746913
+          - 0.18374495208263397
+          - 0.0
+          - 0.47009503841400146
+          - 0.700372040271759
+          std:
+          - 0.28622451424598694
+          - 0.9085981845855713
+          - 0.6900886297225952
+          - 0.5145737528800964
+          - 0.024824900552630424
+          - 0.7737277746200562
+          - 0.0
+          - 0.3019956946372986
+          - 0.9365604519844055
+          - 0.720547616481781
+          - 0.5160077214241028
+          - 0.04248249530792236
+          - 0.4887017011642456
+          - 0.0
+          - 0.4873206615447998
+          - 0.4321308135986328
+          q01:
+          - -0.991405725479126
+          - -1.813583730836399e-05
+          - -3.476843994576484e-05
+          - -1.4938021898269653
+          - -0.04476987197995186
+          - -2.274427652359009
+          - 0.0
+          - -2.3844564566388726e-05
+          - -1.9354150936123915e-05
+          - -3.23197637044359e-05
+          - -1.455243706703186
+          - -0.11395290493965149
+          - -1.3077747821807861
+          - 0.0
+          - 0.0
+          - 0.0
+          q99:
+          - 2.0988931282772683e-05
+          - 2.410663366317749
+          - 2.2642898559570312
+          - 1.6971260265563615e-05
+          - 0.09090737998485565
+          - 1.4203816652297974
+          - 0.0
+          - 1.0001307725906372
+          - 2.56237530708313
+          - 2.397764205932617
+          - 2.7463351216283627e-05
+          - 0.12876388430595398
+          - 1.4705802202224731
+          - 0.0
+          - 1.0
+          - 1.0
+          count:
+          - 277350.0
+          names:
+          - '0'
+          - '1'
+          - '2'
+          - '3'
+          - '4'
+          - '5'
+          - '6'
+          - '7'
+          - '8'
+          - '9'
+          - '10'
+          - '11'
+          - '12'
+          - '13'
+          - '14'
+          - '15'
+          mask:
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+          - true
+        repo_ids:
+        - rmbench_stage2
+    norm_mode: q01_q99
+parallelism:
+  data_parallel_replicate_degree: 1
+  enable_compiled_autograd: false
+  data_parallel_shard_degree: -1
+  fsdp_reshard_after_forward: default
+  context_parallel_config:
+    degree: 1
+    attention_type: ulysses
+    load_balancer: ulysses
+    head_stride: 1
+  tensor_parallel_config:
+    degree: 1
+    enable_async: false
+  data_parallel_config:
+    name: fsdp
+    param_dtype: null
+    reduce_dtype: float32
+    num_replicas: null
+    shard_degree: null
+    wrapping_strategy: full
+    prefetch_factor: 0
+  context_parallel_rotate_method: allgather
+seed: 6198
+epoch: null
+dry_run: false
+ft_llm: true
+ft_vit: false
+ft_connector: true
+ft_action_expert: true
+ft_embedding: lm_head
+optimizer:
+  name: adamw
+  learning_rate: 0.0001
+  weight_decay: 0.01
+  betas:
+  - 0.9
+  - 0.95
+  eps: 1.0e-05
+  connector_learning_rate: 2.5e-05
+  vit_learning_rate: 5.0e-06
+  llm_learning_rate: 2.5e-05
+  frame_selector_learning_rate: 0.0001
+  temporal_token_scorer_learning_rate: 0.0001
+  action_expert_learning_rate: 2.5e-05
+  connector_weight_decay: 0.0
+  vit_weight_decay: 0.0
+  llm_weight_decay: 0.0
+  frame_selector_weight_decay: 0.01
+  temporal_token_scorer_weight_decay: 0.01
+  action_expert_weight_decay: 0.0
+  connector_betas:
+  - 0.9
+  - 0.95
+  vit_betas:
+  - 0.9
+  - 0.95
+  llm_betas:
+  - 0.9
+  - 0.95
+  frame_selector_betas:
+  - 0.9
+  - 0.95
+  temporal_token_scorer_betas:
+  - 0.9
+  - 0.95
+  action_expert_betas:
+  - 0.9
+  - 0.95
+  connector_eps: 1.0e-06
+  vit_eps: 1.0e-06
+  llm_eps: 1.0e-06
+  frame_selector_eps: 1.0e-06
+  temporal_token_scorer_eps: 1.0e-06
+  action_expert_eps: 1.0e-06
+  metrics_log_interval: -1
+scheduler:
+  name: multimodal
+  units: steps
+  t_warmup: 100
+  t_max: null
+  alpha_f: 0.1
+  connector_t_warmup: 200
+  vit_t_warmup: 200
+  llm_t_warmup: 200
+  frame_selector_t_warmup: 200
+  temporal_token_scorer_t_warmup: 200
+  action_expert_t_warmup: 200
+  grad_clip_warmup_steps: null
+  grad_clip_warmup_factor: null
+  warmup_min_lr: 0.0
+data:
+  dataset: null
+  mixture: null
+  root_size_mixture: null
+  kwargs_mixture:
+  - rate: 1.0
+    datasets:
+    - dataset_name: lerobot:rmbench_stage2
+      sampling_rate: null
+      root_size_factor: null
+      message_weight: null
+      override_p_high_res: null
+    name: lerobot:rmbench_stage2
+  split: train
+  seed: 50189
+  pad: null
+  sequence_length: 16384
+  max_text_seq_len: null
+  crop_mode: null
+  shuffle: true
+  start_index: 0
+  packing: null
+  enable_variable_sized_token_pooling: true
+  num_workers: 0
+  drop_last: true
+  pin_memory: false
+  prefetch_factor: 4
+  persistent_workers: false
+  timeout: 0
+  skip_overlong_examples: true
+  skip_missing_vlm_examples: false
+vlm_data: null
+vlm_loader_rate: null
+action_data: null
+action_loader_rate: null
+action_batch_interval: 1
+restore_dataloader: true
+fast_forward_batches: null
+evaluators: []
+eval_interval: 2000
+inf_evaluators: []
+inf_eval_interval: 2000
+save_inloop_predictions: true
+eval_on_last_step: true
+eval_on_load: false
+eval_on: []
+save_folder: outputs/molmoact2-stage2
+reset_parameters: false
+checkpointer_config:
+  save_thread_count: null
+  load_thread_count: null
+  pre_download: false
+  work_dir: null
+  throttle_uploads: false
+canceled_check_interval: 50
+save_interval: 2000
+save_at: null
+save_final_optim: true
+save_num_checkpoints_to_keep: 1
+save_final_unsharded_checkpoint: false
+save_merged_lora_checkpoint: true
+save_interval_ephemeral: null
+save_overwrite: false
+load_path: null
+checkpoint_load_strategy: auto
+reset_optimizer_state: false
+reset_trainer_state: false
+initial_model_checkpoint: allenai/MolmoAct2
+allow_resume: true
+max_duration: 100000
+global_train_batch_size: 4
+device_train_microbatch_size: 1
+max_grad_norm: 1.0
+multi_component_grad_norm: true
+batch_divisor: global_batch
+max_grad_norm_ratio: null
+precision: amp_bf16
+wandb:
+  project: Robocasa
+  entity: ylinwu
+  group: null
+  name: molmoact2-stage2-dinov3-memory-reference-fixed-film-100k-v3
+  tags:
+  - watching
+  log_artifacts: false
+  rank_zero_only: true
+  log_interval: 20
+  allow_resume: true
+  finish_on_sigterm: true
+beaker_log_interval: 50
+speed_monitor:
+  window_size: 20
+  gpu_flops_available: null
+console_log_interval: 20
+enable_timing_logs: false
+gen1_gc_interval: 1
+compile: null
+activation_checkpointing: true
+fsdp:
+  fsdp2: true
+  precision: pure
+  use_orig_params: true
+  wrapping_strategy: null
+  sharding_strategy: FULL_SHARD
+  hybrid_sharding_num_model_replicas: null
+softmax_auxiliary_loss: true
+softmax_auxiliary_loss_scale: 0.0001
+response_logits_only: true
+saliency_score_loss_wt: null
+frame_score_loss_wt: null
+frame_score_loss_type: mse
+frame_score_loss_target: 0.7
+time_limit: null
+extra_steps_after_cancel: 0
+python_profiling: false
+torch_profiling: false
+stop_at: 100000
+stop_after: null
+fused_loss: false
+compile_loss: true
+distributed_eval: true

simulation/rmbench/molmoact2/checkpoint/model.pt

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

simulation/rmbench/molmoact2/manifest.json

@@ -0,0 +1,24 @@
+{
+  "base_model": {
+    "repo": "allenai/MolmoAct2",
+    "revision": "e432d85f6e039edca44afb93c262f3084ab72a9c"
+  },
+  "checkpoint": {
+    "path": "checkpoint/model.pt",
+    "sha256": "ae0dacfcb2a45e82d18266ff5ed49edc26c35c41e038a43065c9a5671773360c"
+  },
+  "checkpoint_config": {
+    "path": "checkpoint/config.yaml",
+    "sha256": "ed4573be786a951ccb3319a090b8443db03acb1a204c827327eb6164ecfb905d"
+  },
+  "checkpoint_step": 100000,
+  "eval_assets_path_prefix": "simulation/rmbench/molmoact2_eval_assets",
+  "eval_assets_repo": "Amazon-FAR/imprint_datasets",
+  "hf_repo": "Amazon-FAR/imprint_models",
+  "memory_enabled": true,
+  "offline_visual_features": false,
+  "path_prefix": "simulation/rmbench/molmoact2",
+  "schema_version": 1,
+  "source_run": "mit0jkx7",
+  "visual_history": "online_wan_and_dinov3"
+}

simulation/rmbench/pi0/README.md

@@ -0,0 +1,34 @@
+# Native Pi0 Memory Policy
+
+This directory contains the native Pi0 IMPRINT checkpoint used for the
+RMBench backbone comparison.
+
+## Configuration
+
+- Policy backbone: native Pi0
+- Memory encoder: full-attention IMPRINT encoder
+- Visual history: pooled DINOv3 features
+- Stage-1 source: `j8t4ebpo`, step 160,000
+- Stage-2 source: `s746sl7k`, step 100,000
+- Memory bottleneck: 256 tokens of width 1024
+- Action horizon: 50
+- Image encoder: frozen in Stage 2
+- Memory encoder: finetuned in Stage 2
+
+The `stage2/step-100000/` directory is a parameter-only Orbax checkpoint.
+Optimizer and training state are intentionally omitted.
+
+Download this prefix with:
+
+```bash
+hf download Amazon-FAR/imprint_models \
+  --include "simulation/rmbench/pi0/**" \
+  --local-dir artifacts/downloads/imprint_models
+```
+
+The matching training store is
+`Amazon-FAR/imprint_datasets/simulation/rmbench/450/`. The released code and
+evaluation harness are on the `pi0` branch of `Amazon-FAR/imprint`.
+
+The checkpoint contains Gemma/PaliGemma-derived weights. The Gemma Terms of
+Use included with the IMPRINT code release apply.

simulation/rmbench/pi0/SHA256SUMS

@@ -0,0 +1,23 @@
+121e80bc7d65582e61c44bf97cbbc7fae2e5245ec97fd62d75e00e90c8658010  README.md
+a1f32c34a9c451cded8682b8e2ac484fdde2df3af806b4f262af8703429bf66e  config/history_config.yaml
+38a89ec49ceb4bd397b0792fe5704441588efa8aaa024bc6c66669c556500091  config/input_provenance.json
+372a92664fadef29d988c33d083c0c2594227ae0078272d80d1d31560086f994  config/norm_stats.json
+b0746fd94e92ebb8ffdc73a102e934202238c4d2d558f0b20ff878230acc409b  manifest.json
+271ec90649ef2f91404db80c799e01e60601a05510f4bb7d9ffcc166e48edce9  stage2/step-100000/_CHECKPOINT_METADATA
+372a92664fadef29d988c33d083c0c2594227ae0078272d80d1d31560086f994  stage2/step-100000/assets/robomme/norm_stats.json
+c52a0aaf801528076bb3c273e81100f81908e8b0b50a1e9b083e6f4a33eafd9d  stage2/step-100000/params/_METADATA
+307790b108b2f99519fda483bfedad812961006d7d0a21c6aa70149beb683326  stage2/step-100000/params/_sharding
+e1ff7698d11457408eb93f14a5bd654c759102c00cceb45648070133409bcd2f  stage2/step-100000/params/array_metadatas/process_0
+85ba1afab930ed68585aab9431d611d73a1c3c09e6703557b018320f80eace4c  stage2/step-100000/params/d/3bd5e7a25312cf2bc1713db9a1d29fa5
+7b71a0eed34c83b3c408232fb342305a3e61df49b44cd84a8c59e71be40cf73e  stage2/step-100000/params/manifest.ocdbt
+5572405112e8e9a60f6755ee127bfb3dc54581e38c9328c1fddb731d89b9b614  stage2/step-100000/params/ocdbt.process_0/d/181a1cb7822e6e4ae74e5d52fd4c3e2c
+6e8e872ea3a0a155be20628b8bb17893de384a3d06de2f924bfec7b4326b9f17  stage2/step-100000/params/ocdbt.process_0/d/24d3ba3b910a02ad89c9e9c5ee773b7d
+b0e00f5770da3b87b76895a636beb69e089d6e96ecb77b3ffcde58de4412cd16  stage2/step-100000/params/ocdbt.process_0/d/3c822efa587f8375db55af5ad6ff4923
+e07b87ad1f9abfd99d86ec5f9b26ca7089caf89b8f859ee3902233018afa736e  stage2/step-100000/params/ocdbt.process_0/d/4371d349f1ebc28e5203786efdbb69cc
+386bb1cf6de7f6a3011524c526bd2c60965f5a2df2e3bcc4522e8d34924a5e82  stage2/step-100000/params/ocdbt.process_0/d/5c003ca5eacb776438fd966f8fe51bdf
+4130dbf0683ef1f0f92ffee27c06d695f8360b0895364228a33591c66fc5f298  stage2/step-100000/params/ocdbt.process_0/d/6f0b57f972221d7c08e1c6e0542b0f59
+034c510b6b4a7b7a01bff53cd076593b5f8e993f6497015f311517b1feced6d7  stage2/step-100000/params/ocdbt.process_0/d/7cf131af294406195d880b45d0b24cdb
+0fcbfee171ad9ec14676b5ed65d48e5fa9f6c88c9eff4edbc260be66e120d754  stage2/step-100000/params/ocdbt.process_0/d/b434ecacfeead5956996de4c9f1c8f19
+3a1cabf9535f848b13bd8013648130093f9937f4c1a3fe95e23ec513bbcd53b9  stage2/step-100000/params/ocdbt.process_0/d/bc0fe1585dbbc9796ccd89ef1cb53d7e
+03e1f1fc55a7493b82d17d583063292ec0e0f1cd8b7d796a13e7cbc8cbb470c4  stage2/step-100000/params/ocdbt.process_0/d/e364328538b1fe570bf0733fb7086ce5
+8bf80ac089f1ca835aedcafd735621ea3c371f10fc19b42efd0a05c371ddced1  stage2/step-100000/params/ocdbt.process_0/manifest.ocdbt

simulation/rmbench/pi0/config/history_config.yaml

@@ -0,0 +1,67 @@
+budget: 256
+num_views: 3
+streaming_obs_horizon: 16
+integration_type: modulation
+memory_token_dim: 1024
+representation_type: cosmo_bottleneck
+memory_feature:
+  img:
+    net: identity
+    input_dim: 768
+  pos:
+    input_dim: 768
+    hidden_dim: 768
+  state:
+    input_dim: 16
+    hidden_dim: 512
+cosmo_bottleneck_memory:
+  visual_feature_stream: dinov3
+  feature_store_path: artifacts/datasets/rmbench/stage1
+  latent_source: offline
+  n_variants: 8
+  num_query_samples: 16
+  variant: Wan2.2-TI2V-5B
+  temporal_factor: 4
+  downsample: 2
+  num_latents: 256
+  num_cross_blocks: 2
+  num_heads: 8
+  cosmos_latent_dim: 48
+  siglip_dim: 768
+  proprio_in_dim: 16
+  proprio_encoding: raw
+  action_in_dim: 16
+  action_encoding: raw
+  lang_emb_dim: 2048
+  action_tokenizer_hidden: 512
+  proprio_tokenizer_hidden: 512
+  pool_h: 2
+  pool_w: 2
+  wan_h: 256
+  wan_w: 256
+  latent_cap: 4096
+  max_recent_tokens: 768
+  max_step_tokens: 1244
+cosmo_decoder:
+  num_heads: 8
+  num_blocks: 2
+  align_dim: 768
+  align_num_patches: 192
+  ffn_inner: 8192
+  action_dim: 16
+  proprio_dim: 16
+  max_query_steps: 2048
+frame_jitter:
+  enabled: false
+  start_jitter: 3
+  stride_choices:
+  - 1
+  - 2
+  - 3
+  interval_jitter: 1
+  canonical_downsample: 2
+stream_aug:
+  enabled: false
+  gauss_sigma: 0.05
+  stream_drop_p: 0.1
+  feat_drop_p: 0.0

simulation/rmbench/pi0/config/input_provenance.json

@@ -0,0 +1,28 @@
+{
+  "action_causality": "strictly_causal",
+  "align_num_patches": 192,
+  "checkpoint_hf_path": "simulation/rmbench/pi0/stage2/step-100000",
+  "checkpoint_hf_repo": "Amazon-FAR/imprint_models",
+  "checkpoint_metadata_sha256": "271ec90649ef2f91404db80c799e01e60601a05510f4bb7d9ffcc166e48edce9",
+  "checkpoint_step": 100000,
+  "eval_consolidate_all": 1,
+  "exact_source_sha256": "ffacf4518beb323c45972975d4b368d953ead306d9389a4a745f3f4d9d0047f6",
+  "history_config_sha256": "e3fd0c5bfd4d7f8f7455e9105a4e2ec0703bdbda67c50c16678288816fa4ba67",
+  "max_recent_tokens": 768,
+  "max_step_tokens": 1244,
+  "memory_language_source": "live",
+  "model_config_name": "mme_wan_stage2_pi0_memft",
+  "model_history_config": "cosmo-wan-bottleneck-modul-rmbench-all9-dinov3-memft.yaml",
+  "model_type": "pi0",
+  "num_latents": 256,
+  "pi05": false,
+  "stage1_checkpoint": "j8t4ebpo_160k",
+  "stage1_checkpoint_hf_path": "simulation/rmbench/full_attention/stage1/stage1_step160000.npz",
+  "stage1_checkpoint_hf_repo": "Amazon-FAR/imprint_models",
+  "stage2_memory_encoder": "finetuned",
+  "training_run": "s746sl7k",
+  "training_store_hf_path": "simulation/rmbench/450",
+  "training_store_hf_repo": "Amazon-FAR/imprint_datasets",
+  "visual_feature_stream": "dinov3",
+  "visual_width": 768
+}

simulation/rmbench/pi0/config/norm_stats.json

@@ -0,0 +1,152 @@
+{
+  "norm_stats": {
+    "actions": {
+      "mean": [
+        -0.2884466350078583,
+        1.116854190826416,
+        0.8176121115684509,
+        -0.6286301016807556,
+        0.0027539930306375027,
+        0.3858006000518799,
+        0.0,
+        0.29012763500213623,
+        1.1434277296066284,
+        0.834911048412323,
+        -0.6164064407348633,
+        0.0033871293999254704,
+        0.1835649162530899,
+        0.0,
+        0.44485998153686523,
+        0.693847119808197
+      ],
+      "std": [
+        0.2924744486808777,
+        0.9123055338859558,
+        0.6921584606170654,
+        0.5167360901832581,
+        0.024822130799293518,
+        0.7730752229690552,
+        0.0,
+        0.3022353947162628,
+        0.9363216161727905,
+        0.7197534441947937,
+        0.5159713625907898,
+        0.042254358530044556,
+        0.4932859241962433,
+        0.0,
+        0.4852803945541382,
+        0.43509986996650696
+      ],
+      "q01": [
+        -1.1566282510757446,
+        -0.00001813583730836399,
+        -0.000035173223295714706,
+        -1.493809461593628,
+        -0.044773634523153305,
+        -2.2744300365448,
+        0.0,
+        -0.000023844564566388726,
+        -0.000019354150936123915,
+        -0.0000323197637044359,
+        -1.4539425373077393,
+        -0.11438962072134018,
+        -1.3077747821807861,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.000020988931282772683,
+        2.410673141479492,
+        2.2643818855285645,
+        0.000017526803276268765,
+        0.09094329178333282,
+        1.4207451343536377,
+        0.0,
+        0.9992868304252625,
+        2.5619561672210693,
+        2.3963286876678467,
+        0.000027463351216283627,
+        0.12538659572601318,
+        1.50703763961792,
+        0.0,
+        1.0,
+        1.0
+      ]
+    },
+    "state": {
+      "mean": [
+        -0.28619107604026794,
+        1.1170974969863892,
+        0.8187518119812012,
+        -0.6292704939842224,
+        0.002742639509961009,
+        0.38321176171302795,
+        0.0,
+        0.287190318107605,
+        1.1249585151672363,
+        0.8228866457939148,
+        -0.6072733998298645,
+        0.003468270180746913,
+        0.18374495208263397,
+        0.0,
+        0.47009503841400146,
+        0.700372040271759
+      ],
+      "std": [
+        0.28622451424598694,
+        0.9085981845855713,
+        0.6900886297225952,
+        0.5145737528800964,
+        0.024824900552630424,
+        0.7737277746200562,
+        0.0,
+        0.3019956946372986,
+        0.9365604519844055,
+        0.720547616481781,
+        0.5160077214241028,
+        0.04248249530792236,
+        0.4887017011642456,
+        0.0,
+        0.4873206615447998,
+        0.4321308135986328
+      ],
+      "q01": [
+        -0.991405725479126,
+        -0.00001813583730836399,
+        -0.00003476843994576484,
+        -1.4938021898269653,
+        -0.04476987197995186,
+        -2.274427652359009,
+        0.0,
+        -0.000023844564566388726,
+        -0.000019354150936123915,
+        -0.0000323197637044359,
+        -1.455243706703186,
+        -0.11395290493965149,
+        -1.3077747821807861,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.000020988931282772683,
+        2.410663366317749,
+        2.2642898559570312,
+        0.000016971260265563615,
+        0.09090737998485565,
+        1.4203816652297974,
+        0.0,
+        1.0001307725906372,
+        2.56237530708313,
+        2.397764205932617,
+        0.000027463351216283627,
+        0.12876388430595398,
+        1.4705802202224731,
+        0.0,
+        1.0,
+        1.0
+      ]
+    }
+  }
+}

simulation/rmbench/pi0/manifest.json

@@ -0,0 +1,26 @@
+{
+  "action_horizon": 50,
+  "architecture": "perceiver_v1",
+  "checkpoint_format": "orbax_params_only",
+  "checkpoint_step": 100000,
+  "hf_repo": "Amazon-FAR/imprint_models",
+  "image_encoder": "frozen",
+  "memory_encoder": "finetuned",
+  "memory_width": 1024,
+  "model_type": "pi0",
+  "num_latents": 256,
+  "path_prefix": "simulation/rmbench/pi0",
+  "payload_bytes": 11534439221,
+  "payload_file_count": 22,
+  "schema_version": 1,
+  "source_eval_bundle": "s746sl7k_100k_pi0_dinov3_memft_causal_c1",
+  "source_run": "s746sl7k",
+  "stage1_checkpoint_step": 160000,
+  "stage1_source_run": "j8t4ebpo",
+  "train_state_included": false,
+  "training_store": {
+    "path_prefix": "simulation/rmbench/450",
+    "repo": "Amazon-FAR/imprint_datasets"
+  },
+  "visual_history": "dinov3"
+}

simulation/rmbench/pi0/stage2/step-100000/_CHECKPOINT_METADATA

@@ -0,0 +1 @@
+{"item_handlers": {"assets": "openpi.training.checkpoints.CallbackHandler", "params": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler", "train_state": "orbax.checkpoint._src.handlers.pytree_checkpoint_handler.PyTreeCheckpointHandler"}, "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1787699801100333867, "commit_timestamp_nsecs": 1787699962354272918, "custom_metadata": {}}

simulation/rmbench/pi0/stage2/step-100000/assets/robomme/norm_stats.json

@@ -0,0 +1,152 @@
+{
+  "norm_stats": {
+    "actions": {
+      "mean": [
+        -0.2884466350078583,
+        1.116854190826416,
+        0.8176121115684509,
+        -0.6286301016807556,
+        0.0027539930306375027,
+        0.3858006000518799,
+        0.0,
+        0.29012763500213623,
+        1.1434277296066284,
+        0.834911048412323,
+        -0.6164064407348633,
+        0.0033871293999254704,
+        0.1835649162530899,
+        0.0,
+        0.44485998153686523,
+        0.693847119808197
+      ],
+      "std": [
+        0.2924744486808777,
+        0.9123055338859558,
+        0.6921584606170654,
+        0.5167360901832581,
+        0.024822130799293518,
+        0.7730752229690552,
+        0.0,
+        0.3022353947162628,
+        0.9363216161727905,
+        0.7197534441947937,
+        0.5159713625907898,
+        0.042254358530044556,
+        0.4932859241962433,
+        0.0,
+        0.4852803945541382,
+        0.43509986996650696
+      ],
+      "q01": [
+        -1.1566282510757446,
+        -0.00001813583730836399,
+        -0.000035173223295714706,
+        -1.493809461593628,
+        -0.044773634523153305,
+        -2.2744300365448,
+        0.0,
+        -0.000023844564566388726,
+        -0.000019354150936123915,
+        -0.0000323197637044359,
+        -1.4539425373077393,
+        -0.11438962072134018,
+        -1.3077747821807861,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.000020988931282772683,
+        2.410673141479492,
+        2.2643818855285645,
+        0.000017526803276268765,
+        0.09094329178333282,
+        1.4207451343536377,
+        0.0,
+        0.9992868304252625,
+        2.5619561672210693,
+        2.3963286876678467,
+        0.000027463351216283627,
+        0.12538659572601318,
+        1.50703763961792,
+        0.0,
+        1.0,
+        1.0
+      ]
+    },
+    "state": {
+      "mean": [
+        -0.28619107604026794,
+        1.1170974969863892,
+        0.8187518119812012,
+        -0.6292704939842224,
+        0.002742639509961009,
+        0.38321176171302795,
+        0.0,
+        0.287190318107605,
+        1.1249585151672363,
+        0.8228866457939148,
+        -0.6072733998298645,
+        0.003468270180746913,
+        0.18374495208263397,
+        0.0,
+        0.47009503841400146,
+        0.700372040271759
+      ],
+      "std": [
+        0.28622451424598694,
+        0.9085981845855713,
+        0.6900886297225952,
+        0.5145737528800964,
+        0.024824900552630424,
+        0.7737277746200562,
+        0.0,
+        0.3019956946372986,
+        0.9365604519844055,
+        0.720547616481781,
+        0.5160077214241028,
+        0.04248249530792236,
+        0.4887017011642456,
+        0.0,
+        0.4873206615447998,
+        0.4321308135986328
+      ],
+      "q01": [
+        -0.991405725479126,
+        -0.00001813583730836399,
+        -0.00003476843994576484,
+        -1.4938021898269653,
+        -0.04476987197995186,
+        -2.274427652359009,
+        0.0,
+        -0.000023844564566388726,
+        -0.000019354150936123915,
+        -0.0000323197637044359,
+        -1.455243706703186,
+        -0.11395290493965149,
+        -1.3077747821807861,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.000020988931282772683,
+        2.410663366317749,
+        2.2642898559570312,
+        0.000016971260265563615,
+        0.09090737998485565,
+        1.4203816652297974,
+        0.0,
+        1.0001307725906372,
+        2.56237530708313,
+        2.397764205932617,
+        0.000027463351216283627,
+        0.12876388430595398,
+        1.4705802202224731,
+        0.0,
+        1.0,
+        1.0
+      ]
+    }
+  }
+}

simulation/rmbench/pi0/stage2/step-100000/params/_METADATA

@@ -0,0 +1 @@
+{"tree_metadata": {"('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoder_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoder_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_0', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_0", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'LayerNorm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "LayerNorm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1152, 1076]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MlpBlock_0', 'Dense_1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MlpBlock_0", "key_type": 2}, {"key": "Dense_1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 1076, 1152]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 16, 72, 288]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 4, 72]}}, "('params', 'PaliGemma', 'img', 'Transformer', 'encoderblock', 'MultiHeadDotProductAttention_0', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "Transformer", "key_type": 2}, {"key": "encoderblock", "key_type": 2}, {"key": "MultiHeadDotProductAttention_0", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [27, 288, 16, 72]}}, "('params', 'PaliGemma', 'img', 'embedding', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [288]}}, "('params', 'PaliGemma', 'img', 'embedding', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "embedding", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [14, 14, 3, 288]}}, "('params', 'PaliGemma', 'img', 'head', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'PaliGemma', 'img', 'head', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "head", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1152, 512]}}, "('params', 'PaliGemma', 'img', 'pos_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "img", "key_type": 2}, {"key": "pos_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1, 64, 1152]}}, "('params', 'PaliGemma', 'llm', 'embedder', 'input_embedding', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "embedder", "key_type": 2}, {"key": "input_embedding", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64288, 2048]}}, "('params', 'PaliGemma', 'llm', 'final_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'PaliGemma', 'llm', 'final_norm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "final_norm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'attn_vec_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "attn_vec_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'kv_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "kv_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 512, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'attn', 'q_einsum_1', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "q_einsum_1", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 8, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'kv_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "kv_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'mem_rms_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "mem_rms_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'out_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "out_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_attn', 'q_einsum_mem', 'w', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_attn", "key_type": 2}, {"key": "q_einsum_mem", "key_type": 2}, {"key": "w", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4, 256, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mem_rms_norm_ffn', 'Dense_0', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mem_rms_norm_ffn", "key_type": 2}, {"key": "Dense_0", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 2048, 4096]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 4096, 2048]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'gating_einsum', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "gating_einsum", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 2, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'mlp_1', 'linear', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "mlp_1", "key_type": 2}, {"key": "linear", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 1024, 1024]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_attention_norm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_attention_norm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 256]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 512]}}, "('params', 'PaliGemma', 'llm', 'layers', 'pre_ffw_norm_1', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "PaliGemma", "key_type": 2}, {"key": "llm", "key_type": 2}, {"key": "layers", "key_type": 2}, {"key": "pre_ffw_norm_1", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [18, 256]}}, "('params', 'action_in_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'action_in_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_in_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 1024]}}, "('params', 'action_out_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8]}}, "('params', 'action_out_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_out_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 32]}}, "('params', 'action_time_mlp_in', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_time_mlp_in", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'action_time_mlp_in', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_time_mlp_in", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512, 1024]}}, "('params', 'action_time_mlp_out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_time_mlp_out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'action_time_mlp_out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "action_time_mlp_out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'ffn', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'ffn', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_kv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_kv', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_q', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_q', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_s', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'norm_s', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '0', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "0", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'ffn', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'ffn', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "ffn", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [1024, 256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_kv', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_kv', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_kv", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_q', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_q', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_q", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_s', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'norm_s', 'scale', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "norm_s", "key_type": 2}, {"key": "scale", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'key', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'key', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "key", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'out', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'out', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "out", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 128, 256]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'query', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'query', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "query", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'value', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2, 128]}}, "('params', 'mem_encoder', 'blocks', '1', 'self_attn', 'value', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "blocks", "key_type": 2}, {"key": "1", "key_type": 2}, {"key": "self_attn", "key_type": 2}, {"key": "value", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256, 8, 128]}}, "('params', 'mem_encoder', 'lang_film', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [512]}}, "('params', 'mem_encoder', 'lang_film', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "lang_film", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [2048, 512]}}, "('params', 'mem_encoder', 'latents', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "latents", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [64, 1024]}}, "('params', 'mem_encoder', 'proj_cosmo', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'proj_cosmo', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_cosmo", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [12, 1024]}}, "('params', 'mem_encoder', 'proj_recent', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'proj_recent', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "proj_recent", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [192, 1024]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_action', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 512]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'tok_action', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_action", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128]}}, "('params', 'mem_encoder', 'tok_proprio', 'l1', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l1", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [4, 512]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'mem_encoder', 'tok_proprio', 'l2', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "mem_encoder", "key_type": 2}, {"key": "tok_proprio", "key_type": 2}, {"key": "l2", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [128, 1024]}}, "('params', 'state_proj', 'bias', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "state_proj", "key_type": 2}, {"key": "bias", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [256]}}, "('params', 'state_proj', 'kernel', 'value')": {"key_metadata": [{"key": "params", "key_type": 2}, {"key": "state_proj", "key_type": 2}, {"key": "kernel", "key_type": 2}, {"key": "value", "key_type": 2}], "value_metadata": {"value_type": "jax.Array", "skip_deserialize": false, "write_shape": [8, 1024]}}}, "use_zarr3": false, "store_array_data_equal_to_fill_value": true, "custom_metadata": null}

simulation/rmbench/pi0/stage2/step-100000/params/_sharding

@@ -0,0 +1 @@
+{"cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2Rlcl9ub3JtLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk11bHRpSGVhZERvdFByb2R1Y3RBdHRlbnRpb25fMC5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLk1scEJsb2NrXzAuRGVuc2VfMS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8wLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuVHJhbnNmb3JtZXIuZW5jb2RlcmJsb2NrLkxheWVyTm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuZW1iZWRkaW5nLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcuaGVhZC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5pbWcucG9zX2VtYmVkZGluZy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZW1iZWRkZXIuaW5wdXRfZW1iZWRkaW5nLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0uZmluYWxfbm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm1lbV9ybXNfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLm91dF9laW5zdW1fbWVtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLmt2X2VpbnN1bV9tZW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9hdHRuLnFfZWluc3VtX21lbS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1lbV9ybXNfbm9ybV9mZm4uRGVuc2VfMC5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5nYXRpbmdfZWluc3VtLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scC5saW5lYXIudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmdhdGluZ19laW5zdW0udmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLm1scF8xLmxpbmVhci52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4uYXR0bl92ZWNfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtLncudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ua3ZfZWluc3VtXzEudy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW0udy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLmF0dG4ucV9laW5zdW1fMS53LnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9hdHRlbnRpb25fbm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLlBhbGlHZW1tYS5sbG0ubGF5ZXJzLnByZV9mZndfbm9ybV8xLnNjYWxlLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fa3YuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fa3Yuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLm5vcm1fcy5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmF0dG4udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmZmbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLmZmbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4wLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fa3YuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fa3Yuc2NhbGUudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcS5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcy5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLm5vcm1fcy5zY2FsZS52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ua2V5LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ua2V5Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ub3V0LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ub3V0Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ucXVlcnkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4ucXVlcnkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4udmFsdWUuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmF0dG4udmFsdWUua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmZmbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLmZmbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi52YWx1ZS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi52YWx1ZS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5rZXkuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5rZXkua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5xdWVyeS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmJsb2Nrcy4xLnNlbGZfYXR0bi5xdWVyeS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null, null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhbmdfZmlsbS5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLmxhdGVudHMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8uYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfY29zbW8ua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50LmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnByb2pfcmVjZW50Lmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDEua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19hY3Rpb24ubDIua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwxLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLm1lbV9lbmNvZGVyLnRva19wcm9wcmlvLmwyLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9pbi5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9pbi5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [\"fsdp\", null], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9vdXQuYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl90aW1lX21scF9vdXQua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [null, \"fsdp\"], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmJpYXMudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9pbl9wcm9qLmtlcm5lbC52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5iaWFzLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLmFjdGlvbl9vdXRfcHJvai5rZXJuZWwudmFsdWU=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnN0YXRlX3Byb2ouYmlhcy52YWx1ZQ==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}","cGFyYW1zLnN0YXRlX3Byb2oua2VybmVsLnZhbHVl":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 4], \"axis_names\": [\"batch\", \"fsdp\"], \"partition_spec\": [], \"device_mesh\": {\"mesh\": [[{\"id\": 0}, {\"id\": 1}, {\"id\": 2}, {\"id\": 3}]]}}"}

simulation/rmbench/pi0/stage2/step-100000/params/array_metadatas/process_0

@@ -0,0 +1 @@
+{"array_metadatas": [{"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoder_norm.scale.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_0.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.LayerNorm_1.scale.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.bias.value", "write_shape": [27, 1076], "chunk_shape": [27, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_0.kernel.value", "write_shape": [27, 1152, 1076], "chunk_shape": [27, 1152, 1076], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MlpBlock_0.Dense_1.kernel.value", "write_shape": [27, 1076, 1152], "chunk_shape": [27, 1076, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.key.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.bias.value", "write_shape": [27, 288], "chunk_shape": [27, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.out.kernel.value", "write_shape": [27, 16, 72, 288], "chunk_shape": [27, 16, 72, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.query.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.bias.value", "write_shape": [27, 4, 72], "chunk_shape": [27, 4, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.Transformer.encoderblock.MultiHeadDotProductAttention_0.value.kernel.value", "write_shape": [27, 288, 16, 72], "chunk_shape": [27, 288, 16, 72], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.bias.value", "write_shape": [288], "chunk_shape": [288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.embedding.kernel.value", "write_shape": [14, 14, 3, 288], "chunk_shape": [14, 14, 3, 288], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.head.kernel.value", "write_shape": [1152, 512], "chunk_shape": [1152, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.img.pos_embedding.value", "write_shape": [1, 64, 1152], "chunk_shape": [1, 64, 1152], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.embedder.input_embedding.value", "write_shape": [64288, 2048], "chunk_shape": [64288, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm.scale.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.final_norm_1.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum.w.value", "write_shape": [18, 8, 256, 512], "chunk_shape": [18, 8, 256, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.attn_vec_einsum_1.w.value", "write_shape": [18, 8, 256, 256], "chunk_shape": [18, 8, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum.w.value", "write_shape": [18, 2, 1, 512, 256], "chunk_shape": [18, 2, 1, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.kv_einsum_1.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum.w.value", "write_shape": [18, 8, 512, 256], "chunk_shape": [18, 8, 512, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.attn.q_einsum_1.w.value", "write_shape": [18, 8, 256, 256], "chunk_shape": [18, 8, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.kv_einsum_mem.w.value", "write_shape": [18, 2, 1, 256, 256], "chunk_shape": [18, 2, 1, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.mem_rms_norm.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.out_einsum_mem.w.value", "write_shape": [18, 4, 256, 256], "chunk_shape": [18, 4, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_attn.q_einsum_mem.w.value", "write_shape": [18, 4, 256, 256], "chunk_shape": [18, 4, 256, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.bias.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mem_rms_norm_ffn.Dense_0.kernel.value", "write_shape": [18, 1024, 512], "chunk_shape": [18, 1024, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.gating_einsum.value", "write_shape": [18, 2, 2048, 4096], "chunk_shape": [18, 2, 2048, 4096], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp.linear.value", "write_shape": [18, 4096, 2048], "chunk_shape": [18, 4096, 2048], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.gating_einsum.value", "write_shape": [18, 2, 1024, 1024], "chunk_shape": [18, 2, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.mlp_1.linear.value", "write_shape": [18, 1024, 1024], "chunk_shape": [18, 1024, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_attention_norm_1.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm.scale.value", "write_shape": [18, 512], "chunk_shape": [18, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.PaliGemma.llm.layers.pre_ffw_norm_1.scale.value", "write_shape": [18, 256], "chunk_shape": [18, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_in_proj.kernel.value", "write_shape": [8, 1024], "chunk_shape": [8, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.bias.value", "write_shape": [8], "chunk_shape": [8], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_out_proj.kernel.value", "write_shape": [256, 32], "chunk_shape": [256, 32], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_time_mlp_in.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_time_mlp_in.kernel.value", "write_shape": [512, 1024], "chunk_shape": [512, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_time_mlp_out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.action_time_mlp_out.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.ffn.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.ffn.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_kv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_kv.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_q.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_q.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_s.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.norm_s.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.0.self_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.ffn.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.ffn.kernel.value", "write_shape": [1024, 256], "chunk_shape": [1024, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_kv.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_kv.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_q.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_q.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_s.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.norm_s.scale.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.key.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.key.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.out.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.out.kernel.value", "write_shape": [8, 128, 256], "chunk_shape": [8, 128, 256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.query.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.query.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.value.bias.value", "write_shape": [2, 128], "chunk_shape": [2, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.blocks.1.self_attn.value.kernel.value", "write_shape": [256, 8, 128], "chunk_shape": [256, 8, 128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.bias.value", "write_shape": [512], "chunk_shape": [512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.lang_film.kernel.value", "write_shape": [2048, 512], "chunk_shape": [2048, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.latents.value", "write_shape": [64, 1024], "chunk_shape": [64, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_cosmo.kernel.value", "write_shape": [12, 1024], "chunk_shape": [12, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.proj_recent.kernel.value", "write_shape": [192, 1024], "chunk_shape": [192, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l1.kernel.value", "write_shape": [4, 512], "chunk_shape": [4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_action.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.bias.value", "write_shape": [128], "chunk_shape": [128], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l1.kernel.value", "write_shape": [4, 512], "chunk_shape": [4, 512], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.mem_encoder.tok_proprio.l2.kernel.value", "write_shape": [128, 1024], "chunk_shape": [128, 1024], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.state_proj.bias.value", "write_shape": [256], "chunk_shape": [256], "ext_metadata": null}}, {"array_metadata": {"param_name": "params.state_proj.kernel.value", "write_shape": [8, 1024], "chunk_shape": [8, 1024], "ext_metadata": null}}]}

simulation/rmbench/pi0/stage2/step-100000/params/d/3bd5e7a25312cf2bc1713db9a1d29fa5

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

simulation/rmbench/pi0/stage2/step-100000/params/ocdbt.process_0/d/181a1cb7822e6e4ae74e5d52fd4c3e2c

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